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__    							= [
    '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__    							= [
    'mlp.dense_4h_to_h.weight',
    'self_attention.dense.weight',
]




def        lowerCAmelCase_			(					__A,  __A							)					->    Union[str, Any]:







      '''simple docstring'''


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




def        lowerCAmelCase_			(					__A							)					->    int:







      '''simple docstring'''


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




def        lowerCAmelCase_			(					__A,  __A,  __A,  __A,  __A							)					->    Union[str, Any]:







      '''simple docstring'''


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

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

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

            UpperCAmelCase__								=  None

            UpperCAmelCase__								=  BloomConfig()

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

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

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

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

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

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

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

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

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

                  UpperCAmelCase__								=  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:
                        UpperCAmelCase__								=  set(other_keys.missing_keys							)
                  else:
                        UpperCAmelCase__								=  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							)
            UpperCAmelCase__								=  pytorch_dump_folder_path + "/" + WEIGHTS_NAME
            UpperCAmelCase__								=  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:
                  UpperCAmelCase__								=  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__":
      UpperCamelCase__    							= 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__    							= 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,
      )




# Copyright 2021 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse

from .config import config_command_parser
from .config_args import default_config_file, load_config_from_file  # noqa: F401
from .default import default_command_parser
from .update import update_command_parser




def        lowerCAmelCase_			(					__A=None							)					->    str:







      '''simple docstring'''


      UpperCAmelCase__								=  argparse.ArgumentParser(add_help=__A,  allow_abbrev=__A							)
      # The main config parser
      UpperCAmelCase__								=  config_command_parser(__A							)
      # The subparser to add commands to
      UpperCAmelCase__								=  config_parser.add_subparsers(title="subcommands",  dest="subcommand"							)

      # Then add other parsers with the parent parser
      default_command_parser(__A,  parents=[parent_parser]							)
      update_command_parser(__A,  parents=[parent_parser]							)

      return config_parser




def        lowerCAmelCase_			(					)					->    Any:







      '''simple docstring'''


      UpperCAmelCase__								=  get_config_parser()
      UpperCAmelCase__								=  config_parser.parse_args()

      if not hasattr(__A,  "func"							):
            config_parser.print_help()
            exit(1							)

      # Run
      args.func(__A							)


if __name__ == "__main__":
      main()



