"""
    #################################################################################################################################################################################################
    # 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 os
from pathlib import Path
import torch
from sapiens_accelerator.utils.modeling import find_tied_parameters
from seamless_communication.models.inference.translator import Translator
from sapiens_transformers import (SeamlessM4TConfig, SeamlessM4TFeatureExtractor, SeamlessM4TModel, SeamlessM4TProcessor, SeamlessM4TTokenizer)
from sapiens_transformers.utils import logging
UNIT_SUPPORTED_LANGUAGES = ["__arb__", "__ben__", "__cat__", "__ces__", "__cmn__", "__cym__", "__dan__", "__deu__", "__eng__", "__est__", "__fin__", "__fra__",
"__hin__", "__ind__", "__ita__", "__jpn__", "__kan__", "__kor__", "__mlt__", "__nld__", "__pes__", "__pol__", "__por__", "__ron__", "__rus__", "__slk__", "__spa__",
"__swe__", "__swh__", "__tam__", "__tel__", "__tgl__", "__tha__", "__tur__", "__ukr__", "__urd__", "__uzn__", "__vie__"]
VOCODER_SUPPORTED_LANGUAGES = ["__arb__", "__ben__", "__cat__", "__ces__", "__cmn__", "__cym__", "__dan__", "__deu__", "__eng__", "__est__", "__fin__", "__fra__",
"__hin__", "__ind__", "__ita__", "__jpn__", "__kor__", "__mlt__", "__nld__", "__pes__", "__pol__", "__por__", "__ron__", "__rus__", "__slk__", "__spa__", "__swe__",
"__swh__", "__tel__", "__tgl__", "__tha__", "__tur__", "__ukr__", "__urd__", "__uzn__", "__vie__"]
MEDIUM_SUPPORTED_LANGUAGES = ["ace","ace_Latn","acm","acq","aeb","afr","ajp","aka","amh","apc","arb","ars","ary","arz","asm","ast","awa","ayr","azb","azj","bak",
"bam","ban","bel","bem","ben","bho","bjn","bjn_Latn","bod","bos","bug","bul","cat","ceb","ces","cjk","ckb","crh","cym","dan","deu","dik","dyu","dzo","ell","eng",
"epo","est","eus","ewe","fao","pes","fij","fin","fon","fra","fur","fuv","gla","gle","glg","grn","guj","hat","hau","heb","hin","hne","hrv","hun","hye","ibo","ilo",
"ind","isl","ita","jav","jpn","kab","kac","kam","kan","kas","kas_Deva","kat","knc","knc_Latn","kaz","kbp","kea","khm","kik","kin","kir","kmb","kon","kor","kmr",
"lao","lvs","lij","lim","lin","lit","lmo","ltg","ltz","lua","lug","luo","lus","mag","mai","mal","mar","min","mkd","plt","mlt","mni","khk","mos","mri","zsm","mya",
"nld","nno","nob","npi","nso","nus","nya","oci","gaz","ory","pag","pan","pap","pol","por","prs","pbt","quy","ron","run","rus","sag","san","sat","scn","shn","sin",
"slk","slv","smo","sna","snd","som","sot","spa","als","srd","srp","ssw","sun","swe","swh","szl","tam","tat","tel","tgk","tgl","tha","tir","taq","taq_Tfng","tpi",
"tsn","tso","tuk","tum","tur","twi","tzm","uig","ukr","umb","urd","uzn","vec","vie","war","wol","xho","ydd","yor","yue","cmn","cmn_Hant","zul",]
LARGE_SUPPORTED_LANGUAGES = ["afr","amh","arb","ary","arz","asm","azj","bel","ben","bos","bul","cat","ceb","ces","ckb","cmn","cmn_Hant","cym","dan","deu","ell",
"eng","est","eus","fin","fra","fuv","gaz","gle","glg","guj","heb","hin","hrv","hun","hye","ibo","ind","isl","ita","jav","jpn","kan","kat","kaz","khk","khm","kir",
"kor","lao","lit","lug","luo","lvs","mai","mal","mar","mkd","mlt","mni","mya","nld","nno","nob","npi","nya","ory","pan","pbt","pes","pol","por","ron","rus","sat",
"slk","slv","sna","snd","som","spa","srp","swe","swh","tam","tel","tgk","tgl","tha","tur","ukr","urd","uzn","vie","yor","yue","zlm","zul"]
def assert_param_count(model_1, model_2):
    count_1 = sum(p[1].numel() for p in model_1.named_parameters() if "final_proj" not in p[0])
    count_2 = sum(p[1].numel() for p in model_2.named_parameters() if "final_proj" not in p[0])
    assert count_1 == count_2, f"{model_1.__class__}: {count_1} != {model_2.__class__}: {count_2}"
def param_count(model): return sum(p[1].numel() for p in model.named_parameters() if "final_proj" not in p[0])
def _grab_best_device(use_gpu=True):
    if torch.cuda.device_count() > 0 and use_gpu: device = "cuda"
    else: device = "cpu"
    return torch.device(device)
logging.set_verbosity_info()
logger = logging.get_logger(__name__)
vocoder_convert_list = [("ups", "hifi_gan.upsampler"), ("conv_pre", "hifi_gan.conv_pre"), ("resblocks", "hifi_gan.resblocks"), ("conv_post", "hifi_gan.conv_post"),
("lang", "language_embedding"), ("spkr", "speaker_embedding"), ("dict.", "unit_embedding."), ("dur_predictor.conv1.0", "dur_predictor.conv1"), ("dur_predictor.conv2.0", "dur_predictor.conv2")]
wav2vec_convert_list = [("speech_encoder_frontend.model_dim_proj", "feature_projection.projection"), ("speech_encoder_frontend.post_extract_layer_norm", "feature_projection.layer_norm"),
("speech_encoder_frontend.pos_encoder.conv", "encoder.pos_conv_embed.conv"), ("speech_encoder.inner.layers", "encoder.layers"), ("speech_encoder.inner_layer_norm", "encoder.layer_norm"),
("speech_encoder.adaptor_layers", "adapter.layers"), ("inner_proj", "intermediate_dense"), ("self_attn.output_proj", "self_attn.linear_out"), ("output_proj", "output_dense"),
("self_attn.k_proj", "self_attn.linear_k"), ("self_attn.v_proj", "self_attn.linear_v"), ("self_attn.q_proj", "self_attn.linear_q"), ("self_attn.sdpa.u_bias", "self_attn.pos_bias_u"),
("self_attn.sdpa.v_bias", "self_attn.pos_bias_v"), ("self_attn.sdpa.r_proj", "self_attn.linear_pos"), ("conv.pointwise_conv1", "conv_module.pointwise_conv1"),
("conv.pointwise_conv2", "conv_module.pointwise_conv2"), ("conv.depthwise_conv", "conv_module.depthwise_conv"), ("conv.batch_norm", "conv_module.batch_norm"),
("conv_layer_norm", "conv_module.layer_norm"), ("speech_encoder.proj1", "intermediate_ffn.intermediate_dense"), ("speech_encoder.proj2", "intermediate_ffn.output_dense"),
("speech_encoder.layer_norm", "inner_layer_norm")]
t2u_convert_list = [("t2u_model.final_proj", "lm_head"), ("t2u_model.", "model."), ("encoder_decoder_attn_layer_norm", "cross_attention_layer_norm"), ("encoder_decoder_attn", "cross_attention"),
("linear_k", "k_proj"), ("linear_v", "v_proj"), ("linear_q", "q_proj"), ("ffn.inner_proj", "ffn.fc1"), ("ffn.output_proj", "ffn.fc2"), ("output_proj", "out_proj"),
("decoder_frontend.embed", "decoder.embed_tokens")]
text_convert_list = [("text_encoder.", ""), ("text_decoder.", ""), ("text_encoder_frontend.embed", "embed_tokens"), ("text_decoder_frontend.embed", "embed_tokens"),
("encoder_decoder_attn_layer_norm", "cross_attention_layer_norm"), ("encoder_decoder_attn", "cross_attention"), ("linear_k", "k_proj"), ("linear_v", "v_proj"),
("linear_q", "q_proj"), ("ffn.inner_proj", "ffn.fc1"), ("ffn.output_proj", "ffn.fc2"), ("output_proj", "out_proj"), ("final_proj", "lm_head")]
CUR_PATH = os.path.dirname(os.path.abspath(__file__))
default_cache_dir = os.path.join(os.path.expanduser("~"), ".cache")
CACHE_DIR = os.path.join(os.getenv("XDG_CACHE_HOME", default_cache_dir), "huggingface", "hub")
def _load_hf_config(model_type="medium"):
    if model_type == "medium":
        kwargs = {'vocab_size': 256206, 't2u_vocab_size': 10082, 'hidden_size': 1024, 'max_position_embeddings': 4096, 'encoder_layers': 12, 'decoder_layers': 12,
        'encoder_ffn_dim': 4096, 'decoder_ffn_dim': 4096, 't2u_encoder_layers': 4, 't2u_decoder_layers': 4, 'speech_encoder_layers': 12}
        return SeamlessM4TConfig(**kwargs)
    else: return SeamlessM4TConfig()
def _convert_model(original_model, hf_model, convert_list, device, unwanted_prefix="model.", filter_state_dict="speech", exclude_state_dict=None):
    state_dict = original_model.state_dict()
    if isinstance(filter_state_dict, str):
        def filter_func(x): return filter_state_dict in x[0]
    else:
        def filter_func(item):
            if exclude_state_dict is not None and exclude_state_dict in item[0]: return False
            for filter_el in filter_state_dict:
                if filter_el in item[0]: return True
            return False
    state_dict = dict(filter(filter_func, state_dict.items()))
    for k, v in list(state_dict.items()):
        new_k = k[len(unwanted_prefix) :]
        for old_layer_name, new_layer_name in convert_list:
            if old_layer_name in new_k: new_k = new_k.replace(old_layer_name, new_layer_name)
        if ".layer_norm" in new_k and new_k.split(".layer_norm")[0][-1].isnumeric(): new_k = new_k.replace("layer_norm", "final_layer_norm")
        state_dict[new_k] = state_dict.pop(k)
    extra_keys = set(state_dict.keys()) - set(hf_model.state_dict().keys())
    extra_keys = set(extra_keys)
    missing_keys = set(hf_model.state_dict().keys()) - set(state_dict.keys())
    missing_keys = set({k for k in missing_keys if "final_logits_bias" not in k})
    if len(extra_keys) != 0: raise ValueError(f"extra keys found: {extra_keys}")
    if len(missing_keys) != 0: raise ValueError(f"missing keys: {missing_keys}")
    hf_model.load_state_dict(state_dict, strict=False)
    n_params = param_count(hf_model)
    logger.info(f"model loaded: {round(n_params/1e6,1)}M params")
    hf_model.eval()
    hf_model.to(device)
    del state_dict
    return hf_model
def load_model(save_dir, model_type, repo_id):
    device = _grab_best_device()
    if model_type == "medium": name = "seamlessM4T_medium"
    else: name = "seamlessM4T_large"
    original_model = Translator(name, "vocoder_36langs", device, torch.float32)
    langs = MEDIUM_SUPPORTED_LANGUAGES if model_type == "medium" else LARGE_SUPPORTED_LANGUAGES
    langs = [f"__{lang}__" for lang in langs]
    vocab_file = os.path.join(os.path.expanduser("~"), "tokenizer", model_type, "tokenizer.model")
    save_dir = os.path.join(save_dir, name)
    Path(save_dir).mkdir(exist_ok=True)
    tokenizer = SeamlessM4TTokenizer(vocab_file, additional_special_tokens=langs)
    sanity_check_lang_id = tokenizer.convert_tokens_to_ids("__fra__")
    tokenizer.save_pretrained(save_dir)
    tokenizer = SeamlessM4TTokenizer.from_pretrained(save_dir)
    if sanity_check_lang_id != tokenizer.convert_tokens_to_ids("__fra__"): raise ValueError(f"Error in tokenizer saving/loading - __fra__ lang id is not coherent: {sanity_check_lang_id} vs {tokenizer.convert_tokens_to_ids('__fra__')}")
    text_decoder_lang_code_to_id = {lang.replace("__", ""): tokenizer.convert_tokens_to_ids(lang) for lang in langs}
    t2u_lang_code_to_id = {code.replace("__", ""): i + 10005 + len(UNIT_SUPPORTED_LANGUAGES) for i, code in enumerate(UNIT_SUPPORTED_LANGUAGES)}
    vocoder_lang_code_to_id = {code.replace("__", ""): i for i, code in enumerate(VOCODER_SUPPORTED_LANGUAGES)}
    fe = SeamlessM4TFeatureExtractor(language_code=langs)
    fe.save_pretrained(save_dir)
    fe = SeamlessM4TFeatureExtractor.from_pretrained(save_dir)
    processor = SeamlessM4TProcessor(feature_extractor=fe, tokenizer=tokenizer)
    processor.save_pretrained(save_dir)
    processor.push_to_hub(repo_id=repo_id, create_pr=True)
    processor = SeamlessM4TProcessor.from_pretrained(save_dir)
    hf_config = _load_hf_config(model_type)
    hf_model = SeamlessM4TModel(hf_config)
    hf_model.generation_config.__setattr__("text_decoder_lang_to_code_id", text_decoder_lang_code_to_id)
    hf_model.generation_config.__setattr__("t2u_lang_code_to_id", t2u_lang_code_to_id)
    hf_model.generation_config.__setattr__("vocoder_lang_code_to_id", vocoder_lang_code_to_id)
    hf_model.vocoder.apply_weight_norm()
    hf_model.vocoder = _convert_model(original_model, hf_model.vocoder, vocoder_convert_list, device, unwanted_prefix="vocoder.code_generator.", filter_state_dict="vocoder")
    hf_model.vocoder.remove_weight_norm()
    wav2vec = hf_model.speech_encoder
    hf_model.speech_encoder = _convert_model(original_model, wav2vec, wav2vec_convert_list, device, unwanted_prefix="model.", filter_state_dict="speech")
    hf_model.t2u_model = _convert_model(original_model, hf_model.t2u_model, t2u_convert_list, device, unwanted_prefix="model.", filter_state_dict="t2u_model")
    hf_model.text_encoder = _convert_model(original_model, hf_model.text_encoder, text_convert_list, device, unwanted_prefix="model.", filter_state_dict=["model.text_encoder"], exclude_state_dict="t2u_model")
    hf_model.text_decoder = _convert_model(original_model, hf_model.text_decoder, text_convert_list, device, unwanted_prefix="model.", filter_state_dict=["model.text_decoder"], exclude_state_dict="t2u_model")
    hf_model.lm_head = _convert_model(original_model, hf_model.lm_head, [("final_proj.", "")], device, unwanted_prefix="model.", filter_state_dict=["model.final_proj"], exclude_state_dict="t2u_model")
    print(find_tied_parameters(hf_model))
    count_1 = param_count(hf_model)
    count_2 = param_count(original_model)
    print(f"HF MODEL:{count_1}, ORIGINAL_MODEL: {count_2}, diff:{count_1 - count_2}")
    print(f"HF MODEL excluding embeddings:{hf_model.num_parameters(exclude_embeddings=True)}")
    del original_model
    hf_model.generation_config._from_model_config = False
    hf_model.save_pretrained(save_dir)
    hf_model.push_to_hub(repo_id=repo_id, create_pr=True)
    hf_model = SeamlessM4TModel.from_pretrained(save_dir)
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--model_type", default="medium", type=str, help="Model type.")
    parser.add_argument("--save_dir", default="/home/ubuntu/weights", type=str, help="Path to the output PyTorch model.")
    parser.add_argument("--repo_id", default="facebook/hf-seamless-m4t-medium", type=str, help="Repo ID.")
    args = parser.parse_args()
    load_model(args.save_dir, args.model_type, args.repo_id)
"""
    #################################################################################################################################################################################################
    # 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.                                                                 #
    #################################################################################################################################################################################################
"""
