import json
from statistics import geometric_mean

import torch
from transformers import (
    GPT2LMHeadModel,
    GPT2TokenizerFast,
    BertForMaskedLM,
    BertTokenizerFast,
    AutoTokenizer,
    TransfoXLTokenizer,
    TransfoXLLMHeadModel,
    T5Tokenizer,
    T5ForConditionalGeneration,
    OpenAIGPTLMHeadModel,
    OpenAIGPTTokenizer,
    XLNetTokenizer,
    XLNetLMHeadModel,
    GPTNeoForCausalLM,
    AutoTokenizer,
    GPTJForCausalLM,
    AutoModelForCausalLM,
)
from transformers import (
    AutoTokenizer,
    BloomForCausalLM,
    OPTForCausalLM,
    AutoModelForSeq2SeqLM,
)
from .data_preprocess.gen_obj_alias_clean_dict import judge_obj_in_vocab
from .data_preprocess.wikidata_get import *
import random
import jsonlines
import os
import os.path as osp
import nltk
import math
from tqdm import tqdm

from minkarr.globs import PROJECT_PATH
from minkarr.utils import download_and_unzip, load_json

try:
    nltk.data.find("stopwords")
except LookupError:
    nltk.download("stopwords")
from nltk.corpus import stopwords

random.seed(1)


def data_filter(probing_facts, rel2alias, sub2alias, obj2alias, verbose=0):
    filtered_facts = dict()
    print("filtering data")
    for fact_id in tqdm(probing_facts.keys(), desc="filtered_facts"):
        fact = probing_facts[fact_id]
        sub_id = fact[0]
        rel_id = fact[1]
        obj_id = fact[2]

        sub_gold_info = get_entity_defaut_alias(sub_id, sub2alias)
        obj_gold_info = get_entity_defaut_alias(obj_id, obj2alias)
        if sub_gold_info == None or obj_gold_info == None:
            if verbose:
                print("no defaut alias for sub or obj")
            continue
        sub_gold = sub_gold_info
        obj_gold = obj_gold_info

        if rel_id not in rel2alias.keys():
            if verbose:
                print("no single text rel")
            continue

        rel_gold_list = rel2alias[rel_id]

        filtered_facts[fact_id] = [
            sub_id,
            rel_id,
            obj_id,
            sub_gold,
            rel_gold_list,
            obj_gold,
        ]
        if sub_id not in sub2alias.keys():
            sub2alias[sub_id] = [sub_gold]
        elif sub_gold not in sub2alias[sub_id]:
            sub2alias[sub_id].append(sub_gold)
        if obj_id not in obj2alias.keys():
            obj2alias[obj_id] = [obj_gold]
        elif obj_gold not in obj2alias[obj_id]:
            obj2alias[obj_id].append(obj_gold)

    return filtered_facts, sub2alias, obj2alias


def gmean(input_x, dim):
    log_x = torch.log(input_x)
    return torch.exp(torch.mean(log_x, dim=dim))


def get_entity_defaut_alias(entity_id, whole_dict):
    if entity_id in whole_dict.keys() and len(whole_dict[entity_id]) > 0:
        return whole_dict[entity_id][0]
    else:
        return None


def condition_prob(
    predict_sent, given_sent, tokenizer, model, device, mode="given beta"
):
    if mode == "given beta":
        # predict_sent = o
        prompt = predict_sent  # s,r
        tgt_len = len(
            tokenizer.encode(" " + predict_sent.replace(given_sent, "").strip())
        )
        encodings = tokenizer(prompt, return_tensors="pt")
        input_ids = encodings.input_ids.to(device)
        target_ids = input_ids.clone()
        target_ids[:, :-tgt_len] = -100
        outputs = model(input_ids, labels=target_ids)
        ppl = torch.exp(outputs.loss)
        return 1 / ppl.item()

    if mode == "given alpha":
        # predict_sent = r,o
        prompt = predict_sent  # s
        tgt_len = len(
            tokenizer.encode(" " + predict_sent.replace(given_sent, "").strip())
        )
        encodings = tokenizer(prompt, return_tensors="pt")
        input_ids = encodings.input_ids.to(device)
        target_ids = input_ids.clone()
        target_ids[:, :-tgt_len] = -100
        outputs = model(input_ids, labels=target_ids)
        ppl = torch.exp(outputs.loss)
        return 1 / ppl.item()


def sentence_prob(sentence, tokenizer, model, device):
    encodings = tokenizer(sentence, return_tensors="pt")
    input_ids = encodings.input_ids.to(device)
    target_ids = input_ids.clone()
    outputs = model(input_ids, labels=target_ids)
    ppl = torch.exp(outputs.loss)
    prob = 1 / ppl.item()
    return prob


def rr_bs_sub_replace(
    tokenizer,
    model,
    device,
    sub_id,
    rel_id,
    obj_id,
    sub2alias,
    rel2alias,
    obj2alias,
    rel2sub2rate,
):
    """
    Given rel sub replacement
    score =
    """

    betas = []
    beta_temps = []
    gammas = []
    alphas = sub2alias[sub_id]
    for r_alias in rel2alias[rel_id]:
        r_alias = r_alias.strip(".").strip()
        if len(r_alias) == 0:
            continue
        for alpha in alphas:
            betas.append(r_alias.replace("[X]", alpha))
        beta_temps.append(r_alias.strip(".").strip())

    for o_alias in obj2alias[obj_id]:
        for beta in betas:
            gammas.append(beta.replace("[Y]", o_alias))

    betas = [beta.replace("[Y]", "").strip() for beta in betas]

    with torch.no_grad():

        # sum_beta [P_M(beta) * \delta(s,r) * sum_gamma P_M(gamma|beta) * \delta(o)]
        # betas represents aliases of s,r
        p_numerator_info = dict()
        p_denominator_info = dict()

        p_numerator = 0
        for beta in betas:
            # P_M(beta)
            p_beta = sentence_prob(beta, tokenizer, model, device)
            p_gamma_sum = 0
            gamma_dict = dict()
            for gamma in gammas:
                # P_M(gamma|beta)
                if beta not in gamma:
                    continue
                p_gamma = condition_prob(
                    gamma, beta, tokenizer, model, device, mode="given beta"
                )
                gamma_dict[gamma] = round(p_gamma, 6)
                p_gamma_sum += p_gamma
            if math.isnan(p_beta) or math.isnan(p_gamma_sum):
                continue
            p_numerator += p_beta * p_gamma_sum

        p_numerator_info[beta] = {
            "p_beta": round(p_beta, 6),
            "p_gamma_sum": round(p_gamma_sum, 6),
            "gamma_dict": gamma_dict,
        }

        p_denominator = 0
        subids_list = [other_sub_id for other_sub_id in rel2sub2rate[rel_id].keys()]
        sample_k = 4
        Q_sampled_sub_idxes = random.sample(
            range(len(subids_list)), min(sample_k, len(subids_list))
        )
        Q_sampled_sub_ids = [subids_list[idx] for idx in Q_sampled_sub_idxes]
        for beta_temp in beta_temps:
            for other_sub_id in Q_sampled_sub_ids:
                for alpha in sub2alias[other_sub_id]:
                    if alpha == None or beta_temp == None:
                        continue
                    p_beta = sentence_prob(
                        beta_temp.replace("[X]", alpha).replace("[Y]", "").strip(),
                        tokenizer,
                        model,
                        device,
                    )
                    p_gamma_sum = 0
                    gamma_dict = dict()
                    for o_alias in obj2alias[obj_id]:
                        gamma = beta_temp.replace("[X]", alpha).replace("[Y]", o_alias)
                        p_gamma = condition_prob(
                            gamma,
                            beta_temp.replace("[X]", alpha).replace("[Y]", "").strip(),
                            tokenizer,
                            model,
                            device,
                            mode="given beta",
                        )
                        p_gamma_sum += p_gamma
                        gamma_dict[gamma] = round(p_gamma, 6)
                    p_denominator += p_beta * p_gamma_sum
        p_denominator = p_denominator / min(sample_k, len(subids_list))
        p_denominator_info[beta] = {
            "p_beta": round(p_beta, 6),
            "p_gamma_sum": round(p_gamma_sum, 6),
            "gamma_dict": gamma_dict,
        }

        if p_denominator == 0:
            return None
        rr_result = {
            "score": p_numerator / p_denominator,
            "p_numerator_score": p_numerator,
            "p_numerator_info": p_numerator_info,
            "p_denominator_score": p_denominator,
            "p_denominator_info": p_denominator_info,
        }
        # print(rr_result)
    return rr_result


def rr_bs_rel_replace(
    tokenizer,
    model,
    device,
    sub_id,
    rel_id,
    obj_id,
    sub2alias,
    rel2alias,
    obj2alias,
    obj2rel2rate,
):
    """
    Given sub, rel replacement
    """
    betas = []
    gammas = []
    alphas = sub2alias[sub_id]
    for r_alias in rel2alias[rel_id]:
        r_alias = r_alias.strip(".").strip()
        if len(r_alias) == 0:
            continue
        for alpha in alphas:
            betas.append(r_alias.replace("[X]", alpha))

    for o_alias in obj2alias[obj_id]:
        for beta in betas:
            gammas.append(beta.replace("[Y]", o_alias))

    betas = [beta.replace("[Y]", "").strip() for beta in betas]

    with torch.no_grad():

        # sum_beta [P_M(beta) * \delta(s,r) * sum_gamma P_M(gamma|beta) * \delta(o)]
        # betas represents aliases of s,r
        p_numerator_info = dict()
        p_denominator_info = dict()

        p_numerator = 0
        for beta in betas:
            # P_M(beta)
            p_beta = sentence_prob(beta, tokenizer, model, device)
            p_gamma_sum = 0
            gamma_dict = dict()
            for gamma in gammas:
                # P_M(gamma|beta)
                if beta not in gamma:
                    continue
                p_gamma = condition_prob(
                    gamma, beta, tokenizer, model, device, mode="given beta"
                )
                gamma_dict[gamma] = round(p_gamma, 6)
                p_gamma_sum += p_gamma
            if math.isnan(p_beta) or math.isnan(p_gamma_sum):
                continue
            p_numerator += p_beta * p_gamma_sum

        p_numerator_info[beta] = {
            "p_beta": round(p_beta, 6),
            "p_gamma_sum": round(p_gamma_sum, 6),
            "gamma_dict": gamma_dict,
        }

        p_denominator = 0
        relids_list = list(obj2rel2rate[obj_id].keys())
        sample_k = 4
        Q_sampled_rel_idxes = random.sample(
            range(len(relids_list)), min(sample_k, len(relids_list))
        )
        Q_sampled_rel_ids = [relids_list[idx] for idx in Q_sampled_rel_idxes]
        for alpha in alphas:
            # P_M(alpha)
            p_alpha = sentence_prob(alpha, tokenizer, model, device)
            p_gamma_sum = 0
            gamma_dict = dict()
            for other_rel_id in Q_sampled_rel_ids:
                # random sample from rel2alias[other_rel_id]
                beta = random.sample(rel2alias[other_rel_id], 1)[0]
                # for beta_text in all_r_gammas[alpha]:
                # P_M(gamma|beta) sum P(Obama was born in Hawaii | Obama) all possible strings
                for o_alias in obj2alias[obj_id]:
                    if beta == None or alpha == None or o_alias == None:
                        continue
                    gamma = beta.replace("[X]", alpha).replace("[Y]", o_alias)
                    p_gamma = condition_prob(
                        gamma, alpha, tokenizer, model, device, mode="given alpha"
                    )
                    p_gamma_sum += p_gamma
                    gamma_dict[gamma] = round(p_gamma, 6)
            if math.isnan(p_alpha) or math.isnan(p_gamma_sum):
                continue
            p_denominator += p_alpha * p_gamma_sum
        p_denominator = p_denominator / (min(sample_k, len(relids_list)))
        p_denominator_info[beta] = {
            "p_alpha": round(p_alpha, 6),
            "p_gamma_sum": round(p_gamma_sum, 6),
            "gamma_dict": gamma_dict,
        }

        if p_denominator == 0:
            return None
        rr_result = {
            "score": p_numerator / p_denominator,
            "p_numerator_score": p_numerator,
            "p_numerator_info": p_numerator_info,
            "p_denominator_score": p_denominator,
            "p_denominator_info": p_denominator_info,
        }
    return rr_result


def build_fact_res(
    fact,
    tokenizer,
    model,
    device,
    sub2alias,
    rel2alias,
    obj2alias,
    obj2rel2rate,
    rel2sub2rate,
):
    sub_id, rel_id, obj_id = fact

    rr_bs_sub_result = rr_bs_sub_replace(
        tokenizer,
        model,
        device,
        sub_id,
        rel_id,
        obj_id,
        sub2alias,
        rel2alias,
        obj2alias,
        rel2sub2rate,
    )
    rr_bs_rel_result = rr_bs_rel_replace(
        tokenizer,
        model,
        device,
        sub_id,
        rel_id,
        obj_id,
        sub2alias,
        rel2alias,
        obj2alias,
        obj2rel2rate,
    )

    if rr_bs_sub_result is None or rr_bs_rel_result is None:
        raise Exception

    fact_res = {
        "rr_bs_sub_result": rr_bs_sub_result,
        "rr_bs_rel_result": rr_bs_rel_result,
    }
    return fact_res


def load_model(model_name, device="cuda"):
    if "gpt2" in model_name:
        model = GPT2LMHeadModel.from_pretrained(model_name).to(device)
        tokenizer = GPT2TokenizerFast.from_pretrained(model_name)
    elif "bloom" in model_name:
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        model = BloomForCausalLM.from_pretrained(model_name).to(device)
    elif "opt" in model_name:
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        model = OPTForCausalLM.from_pretrained(
            model_name,
            torch_dtype=torch.float16,
            device_map="auto",
            low_cpu_mem_usage=True,
        )
    elif "llama" in model_name and "alpaca" not in model_name:
        from transformers import LlamaForCausalLM

        if "65B" or "30B" or "7B" in model_name:
            tokenizer = AutoTokenizer.from_pretrained(
                "YOUR_MODEL_PATH/pretrained_models/LLAMA-hf/{}".format(
                    model_name.split("-")[-1]
                ),
                use_fast=False,
            )
            model = LlamaForCausalLM.from_pretrained(
                "YOUR_MODEL_PATH/pretrained_models/LLAMA-hf/{}".format(
                    model_name.split("-")[-1]
                ),
                torch_dtype=torch.float16,
                device_map="auto",
                low_cpu_mem_usage=True,
            )
        else:
            tokenizer = AutoTokenizer.from_pretrained(
                "YOUR_MODEL_PATH/pretrained_models/LLAMA-hf/{}".format(
                    model_name.split("-")[-1]
                ),
                use_fast=False,
            )
            model = LlamaForCausalLM.from_pretrained(
                "YOUR_MODEL_PATH/pretrained_models/LLAMA-hf/{}".format(
                    model_name.split("-")[-1]
                ),
                torch_dtype=torch.float16,
                low_cpu_mem_usage=True,
            ).to(device)
    elif "alpaca" in model_name:
        from transformers import LlamaForCausalLM

        tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)
        model = LlamaForCausalLM.from_pretrained(
            model_name, torch_dtype=torch.float16, low_cpu_mem_usage=True
        ).to(device)
    elif "vicuna" in model_name:
        from transformers import LlamaForCausalLM

        if "7b" or "13b" in model_name:
            tokenizer = AutoTokenizer.from_pretrained(
                "{}pretrained_models/{}".format(PROJECT_PATH, model_name),
                use_fast=False,
            )
            model = LlamaForCausalLM.from_pretrained(
                "{}pretrained_models/{}".format(PROJECT_PATH, model_name),
                torch_dtype=torch.float16,
                device_map="auto",
                low_cpu_mem_usage=True,
            )
    elif "glm" in model_name:
        tokenizer = AutoTokenizer.from_pretrained(
            "THUDM/glm-10b", trust_remote_code=True
        )
        model = AutoModelForSeq2SeqLM.from_pretrained(
            "THUDM/glm-10b", trust_remote_code=True
        ).to(device)
    elif "moss" in model_name:
        tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
        model = (
            AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True)
            .half()
            .cuda()
        )
    elif "dolly" in model_name:
        tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side="left")
        model = AutoModelForCausalLM.from_pretrained(
            model_name, device_map="auto", torch_dtype=torch.bfloat16
        )
    elif "bert" in model_name:
        model = BertForMaskedLM.from_pretrained(model_name).to(device)
        tokenizer = BertTokenizerFast.from_pretrained(model_name)
    elif "t5" in model_name:
        model = T5ForConditionalGeneration.from_pretrained(model_name).to(device)
        tokenizer = T5Tokenizer.from_pretrained(model_name)
    elif "openai-gpt" in model_name:
        model = OpenAIGPTLMHeadModel.from_pretrained(model_name).to(device)
        tokenizer = OpenAIGPTTokenizer.from_pretrained(model_name)
    elif "xlnet" in model_name:
        model = XLNetLMHeadModel.from_pretrained(model_name).to(device)
        tokenizer = XLNetTokenizer.from_pretrained(model_name)
    elif "gpt-j" in model_name:
        model = GPTJForCausalLM.from_pretrained(model_name).to(device)
        tokenizer = AutoTokenizer.from_pretrained(model_name)

    elif "neo" in model_name:
        model = GPTNeoForCausalLM.from_pretrained(model_name).to(device)
        tokenizer = AutoTokenizer.from_pretrained(model_name)
    else:
        tokenizer = TransfoXLTokenizer.from_pretrained("transfo-xl-wt103")
        model = TransfoXLLMHeadModel.from_pretrained("transfo-xl-wt103").to(device)
    model.eval()
    return model, tokenizer


def setup_data_folder() -> None:
    DATA_URL = (
        "https://zenodo.org/records/13144454/files/data_for_KAssess.zip?download=1"
    )
    print("Setup KaRR...")
    # URL of the zip file to download
    # Ensure the 'data' directory exists
    data_path = osp.join(PROJECT_PATH, "data")
    if not osp.exists(data_path):
        os.makedirs(data_path)
    else:
        print("KaRR setup already DONE!")
        return

    # Call the function with the URL
    download_and_unzip(DATA_URL, data_path, retries=10, backoff_factor=1)
    print("KaRR setup DONE!")


def get_karr(fact_res: dict, thresh: int) -> tuple[float, bool]:
    # I need load_dict
    rr_sub_result = fact_res["rr_bs_sub_result"]
    rr_rel_result = fact_res["rr_bs_rel_result"]
    if float(rr_sub_result["score"]) == 0.0:
        rr_sub_result["score"] = 0.000001
    if float(rr_rel_result["score"]) == 0.0:
        rr_rel_result["score"] = 0.000001
    cur_birr = geometric_mean(
        [float(rr_sub_result["score"]), float(rr_rel_result["score"])]
    )
    does_know = cur_birr > thresh
    return cur_birr, does_know


class KaRR:
    def __init__(self, model, tokenizer, device="cuda", thresh=22) -> None:
        self.model_name = model.config.name_or_path
        self.model, self.tokenizer = model, tokenizer
        self.device = device
        self.model_name_replaced = self.model_name.replace("/", "_")
        self.stopwords = stopwords.words("english")
        self.stopwords.extend(["I", "J", "K", "without"])
        self.thresh = thresh

        (
            self.all_trex,
            self.sub2alias,
            self.rel2alias,
            self.obj2alias,
            self.obj2rel2rate,
            self.rel2sub2rate,
        ) = self._load_data()

        sample_facts_num = len(self.all_trex)
        sample_trex_items_ids = random.sample(
            range(len(self.all_trex)), sample_facts_num
        )
        sample_trex_items = dict()
        for sample_id in sample_trex_items_ids:
            if len(sample_trex_items.keys()) > sample_facts_num:
                break

            fact_id = self.all_trex[sample_id]["fact_id"]
            item = self.all_trex[sample_id]
            if "sub_label" in item.keys():
                # for false facts
                false_trans = {
                    "Atlanta": "Q23556",
                    "English": "Q1860",
                    "Chicago": "Q1297",
                    "France": "Q142",
                    "London": "Q84",
                    "French": "Q150",
                    "medicine": "Q11190",
                    "United": "Q30",
                    "NBC": "Q13974",
                    "RCA": "Q50074604",
                    "ABC": "Q169889",
                    "England": "Q21",
                    "science": "Q336",
                }
                if "#" in false_trans[item["obj_label"]] or "#" in item["subj_label"]:
                    print("the fact is not in the original T_REx dataset")
                    continue
                sample_trex_items[str(fact_id)] = (
                    item["subj_label"],
                    item["relation"],
                    false_trans[item["obj_label"]],
                )
            else:
                if "#" in item["obj_label"] or "#" in item["subj_label"]:
                    print("the fact is not in the original T_REx dataset")
                    continue
                sample_trex_items[str(fact_id)] = (
                    item["subj_label"],
                    item["relation"],
                    item["obj_label"],
                )

        _, self.sub2alias, self.obj2alias = data_filter(
            sample_trex_items, self.rel2alias, self.sub2alias, self.obj2alias
        )
        for rel in self.rel2alias:
            if len(self.rel2alias[rel]) > 0:
                self.rel2alias[rel] = random.sample(
                    self.rel2alias[rel], min(4, len(self.rel2alias[rel]))
                )
        for sub in self.sub2alias:
            if len(self.sub2alias[sub]) > 0:
                self.sub2alias[sub] = random.sample(
                    self.sub2alias[sub], min(4, len(self.sub2alias[sub]))
                )
        for obj in self.obj2alias:
            if len(self.obj2alias[obj]) > 0:
                self.obj2alias[obj] = random.sample(
                    self.obj2alias[obj], min(4, len(self.obj2alias[obj]))
                )

    def compute(self, fact: tuple) -> tuple[float, bool]:
        fact_res = build_fact_res(
            fact,
            self.tokenizer,
            self.model,
            self.device,
            self.sub2alias,
            self.rel2alias,
            self.obj2alias,
            self.obj2rel2rate,
            self.rel2sub2rate,
        )
        return get_karr(fact_res, self.thresh)

    def _load_data(self):
        setup_data_folder()
        print("⏰ Loading data....")
        rootdir = f"{PROJECT_PATH}/data/my_TREx_main_new"
        model_name_replaced = self.model_name_replaced

        all_trex = []
        list_path = os.listdir(rootdir)
        for i in range(0, len(list_path)):
            # Construction path
            path = os.path.join(rootdir, list_path[i])
            with jsonlines.open(path) as reader:
                for obj in reader:
                    all_trex.append(obj)

        with open(
            f"{PROJECT_PATH}/data/cleaned_T_REx/rel2sub2rate.json", "r"
        ) as load_f:
            rel2sub2rate = json.load(load_f)
        if "gpt2" in model_name_replaced:
            vocab_path = (
                f"{PROJECT_PATH}/data/cleaned_T_REx/obj2alias_for_gpt2_vocab.json"
            )
        elif "t5" in model_name_replaced:
            vocab_path = (
                f"{PROJECT_PATH}/data/cleaned_T_REx/obj2alias_for_t5-large_vocab.json"
            )
        elif "bloom" in model_name_replaced:
            vocab_path = f"{PROJECT_PATH}/data/cleaned_T_REx/obj2alias_for_bigscience_bloom-560m_vocab.json"
        elif "opt" in model_name_replaced:
            vocab_path = f"{PROJECT_PATH}/data/cleaned_T_REx/obj2alias_for_facebook_opt-125m_vocab.json"
        elif (
            "llama" in model_name_replaced
            or "alpaca" in model_name_replaced
            or "vicuna" in model_name_replaced
        ):
            vocab_path = f"{PROJECT_PATH}/data/cleaned_T_REx/obj2alias_for_decapoda-research_llama-7b-hf_vocab.json"
        elif "glm" in model_name_replaced or "moss" in model_name_replaced:
            vocab_path = (
                f"{PROJECT_PATH}/data/cleaned_T_REx/obj2alias_for_glm_vocab.json"
            )
        else:
            vocab_path = f"{PROJECT_PATH}/data/cleaned_T_REx/obj2alias_for_{model_name_replaced}_vocab.json"

        if not os.path.exists(vocab_path):
            self._generate_alias()

        obj2alias = load_json(vocab_path)
        if "openai-opt" in model_name_replaced:
            for obj_id in obj2alias.keys():
                cur_obj_aliases = obj2alias[obj_id]
                for obj_alias_id in range(len(cur_obj_aliases)):
                    cur_obj_aliases[obj_alias_id] = cur_obj_aliases[
                        obj_alias_id
                    ].lower()
        sub2alias = load_json(f"{PROJECT_PATH}/data/cleaned_T_REx/allsub2alias.json")
        obj2rel2rate = load_json(f"{PROJECT_PATH}/data/cleaned_T_REx/obj2rel2rate.json")
        # rel alias filter
        rel2alias = load_json(f"{PROJECT_PATH}/data/relation2template.json")
        print("😄 All data loaded.\n")
        return all_trex, sub2alias, rel2alias, obj2alias, obj2rel2rate, rel2sub2rate

    def _generate_alias(self, verbose=0) -> None:
        save_path = f"{PROJECT_PATH}/data/cleaned_T_REx/obj2alias_for_{self.model_name_replaced}_vocab.json"
        if osp.exists(save_path):
            print("Aliases already generated for %s!" % self.model_name)

        with open(
            f"{PROJECT_PATH}/data/cleaned_T_REx/allobj2alias.json", "r"
        ) as load_f:
            obj2alias = json.load(load_f)

        save_dict = {}
        valid_obj_num = 0
        for obj_id in tqdm(obj2alias.keys(), desc="Alias Gen. %s" % self.model_name):
            origial_aliases = obj2alias[obj_id]
            save_dict[obj_id] = []
            for alias in origial_aliases:
                if alias == None:
                    continue
                input_ids = (
                    self.tokenizer(alias, return_tensors="pt")
                    .to(self.device)
                    .input_ids[0]
                )
                if judge_obj_in_vocab(self.tokenizer, alias, input_ids, verbose):
                    save_dict[obj_id].append(alias)
                    valid_obj_num += 1
                elif verbose:
                    print(
                        f"alias: {alias}, judge: {judge_obj_in_vocab(self.tokenizer, alias, input_ids)}"
                    )
        with open(save_path, "w") as write_f:
            json.dump(save_dict, write_f, indent=4, ensure_ascii=False)
        print("Number of aliases per object =", valid_obj_num / len(save_dict.keys()))
