"""
    #################################################################################################################################################################################################
    # This code library is an adaptation of the original Transformers and was designed, developed and programmed by Sapiens Technology®.                                                            #
    # Any alteration and/or disclosure of this code without prior authorization is strictly prohibited and is subject to legal action that will be forwarded by the Sapiens Technology® legal team. #
    # This set of algorithms aims to download, train, fine-tune and/or infer large language models from various sources and slopes.                                                                 #
    #################################################################################################################################################################################################
"""
import argparse
import json
import os
import re
import torch
from sapiens_transformers import BloomConfig, BloomModel
from sapiens_transformers.file_utils import CONFIG_NAME, WEIGHTS_NAME
from sapiens_transformers.utils import logging
logging.set_verbosity_info()
WEIGHTS_TO_AVERAGE_ENDSWITH = ["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"]
WEIGHTS_WITH_ROW_PARALLELISM_CONTAIN = ["mlp.dense_4h_to_h.weight", "self_attention.dense.weight"]
def layer_name_mapping(key, file):
    layer_rename_map = {'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]
    layer_number = int(re.match(r".*layer_(\d*).*", file)[1])
    layer_number -= 3
    return f"h.{layer_number}." + key
def get_dtype_size(dtype):
    if dtype == torch.bool: return 1 / 8
    bit_search = re.search(r"[^\d](\d+)$", str(dtype))
    if bit_search is None: raise ValueError(f"`dtype` is not a valid dtype: {dtype}.")
    bit_size = int(bit_search.groups()[0])
    return bit_size // 8
def convert_bloom_checkpoint_to_pytorch(bloom_checkpoint_path, bloom_config_file, pytorch_dump_folder_path, shard_model, pretraining_tp):
    if bloom_config_file == "": config = BloomConfig()
    else: config = BloomConfig.from_json_file(bloom_config_file)
    if shard_model:
        file_names = os.listdir(bloom_checkpoint_path)
        file_names = sorted(filter(lambda s: s.startswith("layer") and "model_00" in s, file_names))
        index_dict = {"weight_map": {}, "metadata": {}}
        total_size = 0
        missing_keys = None
        config = BloomConfig()
        for j, file in enumerate(file_names):
            print("Processing file: {}".format(file))
            tensors = None
            for i in range(pretraining_tp):
                f_name = file.replace("model_00", f"model_0{i}")
                temp = torch.load(os.path.join(bloom_checkpoint_path, f_name), map_location="cpu")
                keys = list(temp.keys())
                for key in keys: temp[layer_name_mapping(key, file)] = temp.pop(key)
                if tensors is None: tensors = temp
                else:
                    for key in tensors.keys():
                        if any(key.endswith(end) for end in WEIGHTS_TO_AVERAGE_ENDSWITH): tensors[key] += temp[key]
                        else:
                            cat_dim = 1 if any(text in key for text in WEIGHTS_WITH_ROW_PARALLELISM_CONTAIN) else 0
                            tensors[key] = torch.cat([tensors[key], temp[key]], dim=cat_dim)
            for key in tensors.keys():
                if any(key.endswith(end) for end in WEIGHTS_TO_AVERAGE_ENDSWITH): tensors[key] = tensors[key] / pretraining_tp
            torch.save(tensors, os.path.join(pytorch_dump_folder_path, "pytorch_model_{}-of-{}.bin".format(str(j + 1).zfill(5), str(len(file_names)).zfill(5))))
            for key in tensors.keys():
                value = tensors[key]
                total_size += value.numel() * get_dtype_size(value.dtype)
                if key not in index_dict["weight_map"]: index_dict["weight_map"][key] = "pytorch_model_{}-of-{}.bin".format(str(j + 1).zfill(5), str(len(file_names)).zfill(5))
        config = BloomConfig()
        pytorch_config_dump_path = pytorch_dump_folder_path + "/" + CONFIG_NAME
        index_dict["metadata"]["total_size"] = total_size
        with open(pytorch_config_dump_path, "w", encoding="utf-8") as f: f.write(config.to_json_string())
        with open(os.path.join(pytorch_dump_folder_path, WEIGHTS_NAME + ".index.json"), "w", encoding="utf-8") as f:
            json_config = json.dumps(index_dict, indent=2, sort_keys=True) + "\n"
            f.write(json_config)
    else:
        model = BloomModel(config)
        file_names = os.listdir(bloom_checkpoint_path)
        file_names = sorted(filter(lambda s: s.startswith("layer") and "model_00" in s, file_names))
        missing_keys = None
        for i, file in enumerate(file_names):
            tensors = None
            for i in range(pretraining_tp):
                f_name = file.replace("model_00", f"model_0{i}")
                temp = torch.load(os.path.join(bloom_checkpoint_path, f_name), map_location="cpu")
                keys = list(temp.keys())
                for key in keys: temp[layer_name_mapping(key, file)] = temp.pop(key)
                if tensors is None: tensors = temp
                else:
                    for key in tensors.keys():
                        if any(key.endswith(end) for end in WEIGHTS_TO_AVERAGE_ENDSWITH): tensors[key] += temp[key]
                        else:
                            cat_dim = 1 if any(text in key for text in WEIGHTS_WITH_ROW_PARALLELISM_CONTAIN) else 0
                            tensors[key] = torch.cat([tensors[key], temp[key]], dim=cat_dim)
            for key in tensors.keys():
                if any(key.endswith(end) for end in WEIGHTS_TO_AVERAGE_ENDSWITH): tensors[key] = tensors[key] / pretraining_tp
            other_keys = model.load_state_dict(tensors, strict=False)
            assert not other_keys.unexpected_keys, f"The keys {other_keys.unexpected_keys} are unexpected"
            if missing_keys is None: missing_keys = set(other_keys.missing_keys)
            else: missing_keys = missing_keys.intersection(set(other_keys.missing_keys))
        assert not missing_keys, f"The keys {missing_keys} are missing"
        os.makedirs(pytorch_dump_folder_path, exist_ok=True)
        pytorch_weights_dump_path = pytorch_dump_folder_path + "/" + WEIGHTS_NAME
        pytorch_config_dump_path = 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: model = model.to(config.torch_dtype)
        torch.save(model.state_dict(), pytorch_weights_dump_path)
        print(f"Save configuration file to {pytorch_config_dump_path}")
        with open(pytorch_config_dump_path, "w", encoding="utf-8") as f: f.write(config.to_json_string())
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    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. \nThis 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")
    args = 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)
"""
    #################################################################################################################################################################################################
    # This code library is an adaptation of the original Transformers and was designed, developed and programmed by Sapiens Technology®.                                                            #
    # Any alteration and/or disclosure of this code without prior authorization is strictly prohibited and is subject to legal action that will be forwarded by the Sapiens Technology® legal team. #
    # This set of algorithms aims to download, train, fine-tune and/or infer large language models from various sources and slopes.                                                                 #
    #################################################################################################################################################################################################
"""
