"""
    #################################################################################################################################################################################################
    # 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 fnmatch
import re
import torch
from sapiens_transformers import (DacConfig, DacFeatureExtractor, DacModel, logging)
logging.set_verbosity_info()
logger = logging.get_logger("sapiens_transformers.models.dac")
def match_pattern(string, pattern):
    pattern_parts = pattern.split(".")
    string_parts = string.split(".")
    pattern_block_count = string_block_count = 0
    for part in pattern_parts:
        if part.startswith("block"): pattern_block_count += 1
    for part in string_parts:
        if part.startswith("block"): string_block_count += 1
    return fnmatch.fnmatch(string, pattern) and string_block_count == pattern_block_count
TOP_LEVEL_KEYS = []
IGNORE_KEYS = []
MAPPING_ENCODER = {"encoder.block.0": ["encoder.conv1"], "encoder.block.5": ["encoder.snake1"], "encoder.block.6": ["encoder.conv2"],
"encoder.block.*.block.*.block.0".replace("*", r"\d+"): ["encoder.block", "res_unit", "snake1"], "encoder.block.*.block.*.block.1".replace("*", r"\d+"): ["encoder.block", "res_unit", "conv1"],
"encoder.block.*.block.*.block.2".replace("*", r"\d+"): ["encoder.block", "res_unit", "snake2"], "encoder.block.*.block.*.block.3".replace("*", r"\d+"): ["encoder.block", "res_unit", "conv2"],
"encoder.block.*.block.3".replace("*", r"\d+"): ["encoder.block", "snake1"], "encoder.block.*.block.4".replace("*", r"\d+"): ["encoder.block", "conv1"]}
MAPPING_QUANTIZER = {'quantizer.quantizers.*': ['quantizer.quantizers.*']}
MAPPING_DECODER = {"decoder.model.0": ["decoder.conv1"], "decoder.model.5": ["decoder.snake1"], "decoder.model.6": ["decoder.conv2"], "decoder.model.*.block.0".replace("*", r"\d+"): ["decoder.block", "snake1"],
"decoder.model.*.block.1".replace("*", r"\d+"): ["decoder.block", "conv_t1"], "decoder.model.*.block.*.block.0".replace("*", r"\d+"): ["decoder.block", "res_unit", "snake1"], "decoder.model.*.block.*.block.1".replace("*", r"\d+"): ["decoder.block", "res_unit", "conv1"],
"decoder.model.*.block.*.block.2".replace("*", r"\d+"): ["decoder.block", "res_unit", "snake2"], "decoder.model.*.block.*.block.3".replace("*", r"\d+"): ["decoder.block", "res_unit", "conv2"]}
MAPPING = {**MAPPING_ENCODER, **MAPPING_QUANTIZER, **MAPPING_DECODER}
def set_recursively(hf_pointer, key, value, full_name, weight_type):
    for attribute in key.split("."): hf_pointer = getattr(hf_pointer, attribute)
    if weight_type is not None: hf_shape = getattr(hf_pointer, weight_type).shape
    else: hf_shape = hf_pointer.shape
    if hf_shape != value.shape: raise ValueError(f"Shape of hf {key + '.' + weight_type if weight_type is not None else ''} is {hf_shape}, but should be {value.shape} for {full_name}")
    if weight_type == "weight": hf_pointer.weight.data = value
    elif weight_type == "weight_g": hf_pointer.weight_g.data = value
    elif weight_type == "weight_v": hf_pointer.weight_v.data = value
    elif weight_type == "bias": hf_pointer.bias.data = value
    elif weight_type == "alpha": hf_pointer.alpha.data = value
    logger.info(f"{key + ('.' + weight_type if weight_type is not None else '')} was initialized from {full_name}.")
def should_ignore(name, ignore_keys):
    for key in ignore_keys:
        if key.endswith(".*"):
            if name.startswith(key[:-1]): return True
        elif ".*." in key:
            prefix, suffix = key.split(".*.")
            if prefix in name and suffix in name: return True
        elif key in name: return True
    return False
def recursively_load_weights(orig_dict, hf_model, model_name):
    unused_weights = []
    if model_name not in ["dac_16khz", "dac_24khz", "dac_44khz"]: raise ValueError(f"Unsupported model: {model_name}")
    for name, value in orig_dict.items():
        is_used = False
        for key, mapped_key in MAPPING.items():
            regex = re.compile(key)
            if regex.search(name):
                if len(mapped_key) == 1:
                    if mapped_key[0][0] == "q": mapped_key = ".".join(name.split(".")[:-1])
                    else: mapped_key = mapped_key[0]
                elif len(mapped_key) == 3:
                    integers = re.findall(r"\b\d+\b", name)
                    if mapped_key[0][0] == "d": mapped_key = "{}.{}.{}{}.{}".format(mapped_key[0], str(int(integers[0]) - 1), mapped_key[1], str(int(integers[1]) - 1), mapped_key[2])
                    else: mapped_key = "{}.{}.{}{}.{}".format(mapped_key[0], str(int(integers[0]) - 1), mapped_key[1], str(int(integers[1]) + 1), mapped_key[2])
                elif len(mapped_key) == 2:
                    integers = re.findall(r"\b\d+\b", name)
                    mapped_key = "{}.{}.{}".format(mapped_key[0], str(int(integers[0]) - 1), mapped_key[1])
                is_used = True
                if "weight_g" in name: weight_type = "weight_g"
                elif "weight_v" in name: weight_type = "weight_v"
                elif "bias" in name: weight_type = "bias"
                elif "alpha" in name: weight_type = "alpha"
                elif "weight" in name: weight_type = "weight"
                set_recursively(hf_model, mapped_key, value, name, weight_type)
        if not is_used: unused_weights.append(name)
    print(list(set(unused_weights)))
    logger.warning(f"Unused weights: {unused_weights}")
@torch.no_grad()
def convert_checkpoint(model_name, checkpoint_path, pytorch_dump_folder_path, sample_rate=16000, repo_id=None):
    model_dict = torch.load(checkpoint_path, "cpu")
    config = DacConfig()
    metadata = model_dict["metadata"]["kwargs"]
    config.encoder_hidden_size = metadata["encoder_dim"]
    config.downsampling_ratios = metadata["encoder_rates"]
    config.codebook_size = metadata["codebook_size"]
    config.n_codebooks = metadata["n_codebooks"]
    config.codebook_dim = metadata["codebook_dim"]
    config.decoder_hidden_size = metadata["decoder_dim"]
    config.upsampling_ratios = metadata["decoder_rates"]
    config.quantizer_dropout = float(metadata["quantizer_dropout"])
    config.sampling_rate = sample_rate
    model = DacModel(config)
    feature_extractor = DacFeatureExtractor()
    feature_extractor.sampling_rate = sample_rate
    original_checkpoint = model_dict["state_dict"]
    model.apply_weight_norm()
    recursively_load_weights(original_checkpoint, model, model_name)
    model.remove_weight_norm()
    model.save_pretrained(pytorch_dump_folder_path)
    if repo_id:
        print("Pushing to the hub...")
        feature_extractor.push_to_hub(repo_id)
        model.push_to_hub(repo_id)
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", default="dac_44khz", type=str, help="The model to convert. Should be one of 'dac_16khz', 'dac_24khz', 'dac_44khz'.")
    parser.add_argument("--checkpoint_path", required=True, default=None, type=str, help="Path to original checkpoint")
    parser.add_argument("--pytorch_dump_folder_path", required=True, default=None, type=str, help="Path to the output PyTorch model.")
    parser.add_argument("--push_to_hub", default=None, type=str, help="Where to upload the converted model on the HF hub.")
    parser.add_argument("--sample_rate", default=None, type=str, help="Sample rate used by DacFeatureExtractor")
    args = parser.parse_args()
    convert_checkpoint(args.model, args.checkpoint_path, args.pytorch_dump_folder_path, args.sample_rate, args.push_to_hub)
"""
    #################################################################################################################################################################################################
    # 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.                                                                 #
    #################################################################################################################################################################################################
"""
