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      :			Optional[int]							       =		[
    '''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      :			Any							       =		[
    '''mlp.dense_4h_to_h.weight''',
    '''self_attention.dense.weight''',
]



def      __SCREAMING_SNAKE_CASE  ( __UpperCamelCase   :	Tuple			,					__UpperCamelCase   :	Dict     )							-> Optional[Any]:







  """simple docstring"""





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



def      __SCREAMING_SNAKE_CASE  ( __UpperCamelCase   :	Union[str, Any]     )							-> str:







  """simple docstring"""





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



def      __SCREAMING_SNAKE_CASE  ( __UpperCamelCase   :	Any			,					__UpperCamelCase   :	str			,					__UpperCamelCase   :	int			,					__UpperCamelCase   :	List[str]			,					__UpperCamelCase   :	int     )							-> Dict:







  """simple docstring"""





  if bloom_config_file == "":
    SCREAMING_SNAKE_CASE__      							=				BloomConfig()
  else:
    SCREAMING_SNAKE_CASE__      							=				BloomConfig.from_json_file(__UpperCamelCase     )

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

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

    SCREAMING_SNAKE_CASE__      							=				None

    SCREAMING_SNAKE_CASE__      							=				BloomConfig()

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

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

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

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

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

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

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

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

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

      SCREAMING_SNAKE_CASE__      							=				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:
        SCREAMING_SNAKE_CASE__      							=				set(other_keys.missing_keys     )
      else:
        SCREAMING_SNAKE_CASE__      							=				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     )
    SCREAMING_SNAKE_CASE__      							=				pytorch_dump_folder_path + """/""" + WEIGHTS_NAME
    SCREAMING_SNAKE_CASE__      							=				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:
      SCREAMING_SNAKE_CASE__      							=				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__":
   __lowerCamelCase      :			Dict							       =		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      :			Dict							       =		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 .lib import (
    Matrix,
    Vector,
    axpy,
    square_zero_matrix,
    unit_basis_vector,
    zero_vector,
)







class __snake_case    (						unittest.TestCase		):




   def      __a						(   self      :    Dict				):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__      							=				Vector([1, 2, 3]				)
     self.assertEqual(x.component(0				)	,     1				)
     self.assertEqual(x.component(2				)	,     3				)
     SCREAMING_SNAKE_CASE__      							=				Vector()




   def      __a						(   self      :    int				):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__      							=				Vector([0, 0, 0, 0, 0, 1]				)
     self.assertEqual(str(_lowercase				)	,     """(0,0,0,0,0,1)"""				)




   def      __a						(   self      :    Tuple				):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__      							=				Vector([1, 2, 3, 4]				)
     self.assertEqual(len(_lowercase				)	,     4				)




   def      __a						(   self      :    Optional[Any]				):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__      							=				Vector([1, 2]				)
     SCREAMING_SNAKE_CASE__      							=				Vector([1, 2, 3, 4, 5]				)
     SCREAMING_SNAKE_CASE__      							=				Vector([0, 0, 0, 0, 0, 0, 0, 0, 0, 0]				)
     SCREAMING_SNAKE_CASE__      							=				Vector([1, -1, 1, -1, 2, -3, 4, -5]				)
     self.assertAlmostEqual(x.euclidean_length()	,     2.2_36	,     3				)
     self.assertAlmostEqual(y.euclidean_length()	,     7.4_16	,     3				)
     self.assertEqual(z.euclidean_length()	,     0				)
     self.assertAlmostEqual(w.euclidean_length()	,     7.6_16	,     3				)




   def      __a						(   self      :    Optional[int]				):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__      							=				Vector([1, 2, 3]				)
     SCREAMING_SNAKE_CASE__      							=				Vector([1, 1, 1]				)
     self.assertEqual((x + y).component(0				)	,     2				)
     self.assertEqual((x + y).component(1				)	,     3				)
     self.assertEqual((x + y).component(2				)	,     4				)




   def      __a						(   self      :    Any				):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__      							=				Vector([1, 2, 3]				)
     SCREAMING_SNAKE_CASE__      							=				Vector([1, 1, 1]				)
     self.assertEqual((x - y).component(0				)	,     0				)
     self.assertEqual((x - y).component(1				)	,     1				)
     self.assertEqual((x - y).component(2				)	,     2				)




   def      __a						(   self      :    str				):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__      							=				Vector([1, 2, 3]				)
     SCREAMING_SNAKE_CASE__      							=				Vector([2, -1, 4]				)  # for test of dot product
     SCREAMING_SNAKE_CASE__      							=				Vector([1, -2, -1]				)
     self.assertEqual(str(x * 3.0				)	,     """(3.0,6.0,9.0)"""				)
     self.assertEqual((a * b)	,     0				)




   def      __a						(   self      :    Union[str, Any]				):



     """simple docstring"""






     self.assertEqual(str(zero_vector(10				)				).count("""0"""				)	,     10				)




   def      __a						(   self      :    str				):



     """simple docstring"""






     self.assertEqual(str(unit_basis_vector(3	,     1				)				)	,     """(0,1,0)"""				)




   def      __a						(   self      :    Any				):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__      							=				Vector([1, 2, 3]				)
     SCREAMING_SNAKE_CASE__      							=				Vector([1, 0, 1]				)
     self.assertEqual(str(axpy(2	,     _lowercase	,     _lowercase				)				)	,     """(3,4,7)"""				)




   def      __a						(   self      :    str				):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__      							=				Vector([1, 0, 0, 0, 0, 0]				)
     SCREAMING_SNAKE_CASE__      							=				x.copy()
     self.assertEqual(str(_lowercase				)	,     str(_lowercase				)				)




   def      __a						(   self      :    str				):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__      							=				Vector([1, 0, 0]				)
     x.change_component(0	,     0				)
     x.change_component(1	,     1				)
     self.assertEqual(str(_lowercase				)	,     """(0,1,0)"""				)




   def      __a						(   self      :    Optional[Any]				):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__      							=				Matrix([[1, 2, 3], [2, 4, 5], [6, 7, 8]]	,     3	,     3				)
     self.assertEqual("""|1,2,3|\n|2,4,5|\n|6,7,8|\n"""	,     str(_lowercase				)				)




   def      __a						(   self      :    List[Any]				):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__      							=				Matrix([[1, 2, 3], [2, 4, 5], [6, 7, 8]]	,     3	,     3				)
     SCREAMING_SNAKE_CASE__      							=				[[-3, -14, -10], [-5, -10, -5], [-2, -1, 0]]
     for x in range(a.height()				):
       for y in range(a.width()				):
         self.assertEqual(minors[x][y]	,     a.minor(_lowercase	,     _lowercase				)				)




   def      __a						(   self      :    Dict				):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__      							=				Matrix([[1, 2, 3], [2, 4, 5], [6, 7, 8]]	,     3	,     3				)
     SCREAMING_SNAKE_CASE__      							=				[[-3, 14, -10], [5, -10, 5], [-2, 1, 0]]
     for x in range(a.height()				):
       for y in range(a.width()				):
         self.assertEqual(cofactors[x][y]	,     a.cofactor(_lowercase	,     _lowercase				)				)




   def      __a						(   self      :    Union[str, Any]				):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__      							=				Matrix([[1, 2, 3], [2, 4, 5], [6, 7, 8]]	,     3	,     3				)
     self.assertEqual(-5	,     a.determinant()				)




   def      __a						(   self      :    List[Any]				):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__      							=				Matrix([[1, 2, 3], [4, 5, 6], [7, 8, 9]]	,     3	,     3				)
     SCREAMING_SNAKE_CASE__      							=				Vector([1, 2, 3]				)
     self.assertEqual("""(14,32,50)"""	,     str(a * x				)				)
     self.assertEqual("""|2,4,6|\n|8,10,12|\n|14,16,18|\n"""	,     str(a * 2				)				)




   def      __a						(   self      :    Optional[Any]				):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__      							=				Matrix([[1, 2, 3], [2, 4, 5], [6, 7, 8]]	,     3	,     3				)
     a.change_component(0	,     2	,     5				)
     self.assertEqual("""|1,2,5|\n|2,4,5|\n|6,7,8|\n"""	,     str(_lowercase				)				)




   def      __a						(   self      :    Union[str, Any]				):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__      							=				Matrix([[1, 2, 3], [2, 4, 5], [6, 7, 8]]	,     3	,     3				)
     self.assertEqual(7	,     a.component(2	,     1				)	,     0.01				)




   def      __a						(   self      :    List[Any]				):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__      							=				Matrix([[1, 2, 3], [2, 4, 5], [6, 7, 8]]	,     3	,     3				)
     SCREAMING_SNAKE_CASE__      							=				Matrix([[1, 2, 7], [2, 4, 5], [6, 7, 10]]	,     3	,     3				)
     self.assertEqual("""|2,4,10|\n|4,8,10|\n|12,14,18|\n"""	,     str(a + b				)				)




   def      __a						(   self      :    Optional[Any]				):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__      							=				Matrix([[1, 2, 3], [2, 4, 5], [6, 7, 8]]	,     3	,     3				)
     SCREAMING_SNAKE_CASE__      							=				Matrix([[1, 2, 7], [2, 4, 5], [6, 7, 10]]	,     3	,     3				)
     self.assertEqual("""|0,0,-4|\n|0,0,0|\n|0,0,-2|\n"""	,     str(a - b				)				)




   def      __a						(   self      :    Any				):



     """simple docstring"""






     self.assertEqual(
         """|0,0,0,0,0|\n|0,0,0,0,0|\n|0,0,0,0,0|\n|0,0,0,0,0|\n|0,0,0,0,0|\n"""	,     str(square_zero_matrix(5				)				)	,     )


if __name__ == "__main__":
   unittest.main()
