A fundational graph learning framework that solves cross-domain/cross-task classification problems using one model.
import os
import pandas as pd
import torch
import torch_geometric as pyg
from data.ofa_data import OFAPygDataset
from ogb.nodeproppred import PygNodePropPredDataset


def get_data(path):
    data = torch.load(path)
    text = data.raw_texts
    label_names = data.label_names
    nx_g = pyg.utils.to_networkx(data, to_undirected=True)
    edge_index = torch.tensor(list(nx_g.edges())).T
    print(edge_index.size())
    data_dict = data.to_dict()
    data_dict["edge_index"] = edge_index
    new_data = pyg.data.data.Data(**data_dict)
    category_desc = pd.read_csv(
        os.path.join(os.path.dirname(__file__), "categories.csv"), sep=","
    ).values
    ordered_desc = []
    for i, label in enumerate(label_names):
        true_ind = label == category_desc[:, 0]
        ordered_desc.append((label, category_desc[:, 1][true_ind]))
    clean_text = ["feature node. paper title and abstract: " + t for t in text]
    label_text = [
        "prompt node. literature category and description: "
        + desc[0]
        + "."
        + desc[1][0]
        for desc in ordered_desc
    ]
    edge_label_text = [
        "prompt node. two papers have co-citation",
        "prompt node. two papers do not have co-citation",
    ]
    edge_text = [
        "feature edge. connected papers are cited together by other papers."
    ]
    prompt_node_edge_text = [
        "prompt node. link prediction on the papers that are cited together"
    ]
    prompt_node_text = [
        "prompt node. node classification on the paper's category"
    ]
    prompt_edge_text = ["prompt edge."]
    return (
        [new_data],
        [
            clean_text,
            label_text,
            edge_text,
            prompt_node_edge_text,
            prompt_node_text,
            prompt_edge_text,
            edge_label_text,
        ],
        None,
    )


class CoraOFADataset(OFAPygDataset):
    def gen_data(self):
        cur_path = os.path.dirname(__file__)
        return get_data(os.path.join(cur_path, "cora.pt"))

    def add_text_emb(self, data_list, text_emb):
        data_list[0].x_text_feat = text_emb[0]
        data_list[0].label_text_feat = text_emb[1]
        data_list[0].edge_text_feat = text_emb[2]
        data_list[0].prompt_text_edge_feat = text_emb[3]
        data_list[0].prompt_text_feat = text_emb[4]
        data_list[0].prompt_edge_feat = text_emb[5]
        data_list[0].edge_label_feat = text_emb[6]
        return self.collate(data_list)

data/Cora/gen_data.py
LechengKong/OneForAll
import os.path as osp
import torch_geometric as pyg
import torch
import json
from data.ofa_data import OFAPygDataset

entity_lst = []
text_lst = []
with open(osp.join(osp.dirname(__file__), "entity2wikidata.json"), "r") as f:
    data = json.load(f)

for k in data:
    # print(data[k])
    entity_lst.append(k)
    text_lst.append(
        "entity nammes: "
        + data[k]["label"]
        + ", entity alternatives: "
        + ", ".join(data[k]["alternatives"])
        + ". entity descriptions:"
        + data[k]["description"]
        if data[k]["description"] is not None
        else "None"
    )


entity2id = {entity: i for i, entity in enumerate(entity_lst)}


def read_knowledge_graph(files):
    relation2id = {}

    converted_triplets = {}
    rel_list = []
    rel = len(relation2id)

    for file_type, file_path in files.items():

        edges = []
        edge_types = []
        with open(file_path) as f:
            file_data = [line.split() for line in f.read().split("\n")[:-1]]
        unknown_entity = 0
        for triplet in file_data:
            if triplet[0] not in entity2id:
                text_lst.append("entity names: Unknown")
                entity_lst.append(triplet[0])
                entity2id[triplet[0]] = len(entity2id)
                unknown_entity += 1
            if triplet[2] not in entity2id:
                text_lst.append("entity names: Unknown")
                entity_lst.append(triplet[2])
                entity2id[triplet[2]] = len(entity2id)
                unknown_entity += 1
            if triplet[1] not in relation2id:
                relation2id[triplet[1]] = rel
                rel_list.append(triplet[1])
                rel += 1
            edges.append(
                [
                    entity2id[triplet[0]],
                    entity2id[triplet[2]],
                ]
            )
            edge_types.append(relation2id[triplet[1]])
        print(unknown_entity)
        converted_triplets[file_type] = [edges, edge_types]

    new_data = pyg.data.data.Data(
        x=torch.zeros([len(text_lst), 1]),
        edge_index=torch.tensor(converted_triplets["train"][0]).T,
        edge_types=torch.tensor(converted_triplets["train"][1]),
    )

    node_text = [
        "feature node. entity and entity description: " + ent
        for ent in text_lst
    ]
    edge_text = [
        "feature edge. relation between two entities. " + relation
        for relation in rel_list
    ] + [
        "feature edge. relation between two entities. the inverse relation of "
        + relation
        for relation in rel_list
    ]

    prompt_edge_text = ["prompt edge"]
    prompt_node_text = [
        "prompt node. relation type prediction between the connected entities."
    ]
    label_text = [
        "prompt node. relation between two entities. " + relation
        for relation in rel_list
    ]

    return (
        [new_data],
        [
            node_text,
            edge_text,
            label_text,
            prompt_edge_text,
            prompt_node_text,
        ],
        [converted_triplets, rel_list],
    )


class FB15K237OFADataset(OFAPygDataset):
    def gen_data(self):
        cur_path = osp.dirname(__file__)
        names = ["train", "valid", "test"]
        name_dict = {n: osp.join(cur_path, n + ".txt") for n in names}
        return read_knowledge_graph(name_dict)

    def add_text_emb(self, data_list, text_emb):
        data_list[0].x_text_feat = text_emb[0]
        data_list[0].edge_text_feat = text_emb[1]
        data_list[0].edge_label_feat = text_emb[2]
        data_list[0].prompt_edge_feat = text_emb[3]
        data_list[0].prompt_text_feat = text_emb[4]
        return self.collate(data_list)

    def get_idx_split(self):
        return self.side_data[0]

data/FB15K237/gen_data.py
LechengKong/OneForAll
import os
import pandas as pd
import torch
import torch_geometric as pyg
from data.ofa_data import OFAPygDataset
from ogb.nodeproppred import PygNodePropPredDataset


def get_data(path):
    data = torch.load(path)
    print(data)
    text = data.raw_texts
    nx_g = pyg.utils.to_networkx(data, to_undirected=True)
    edge_index = torch.tensor(list(nx_g.edges())).T
    print(edge_index.size())
    data_dict = data.to_dict()
    data_dict["edge_index"] = edge_index
    new_data = pyg.data.data.Data(**data_dict)
    with open(
        os.path.join(os.path.dirname(__file__), "categories.csv"), "r"
    ) as f:
        ordered_desc = f.read().split("\n")
    clean_text = ["feature node. paper title and abstract: " + t for t in text]
    label_text = [
        "prompt node. literature category and description: " + desc
        for desc in ordered_desc
    ]
    edge_label_text = [
        "prompt node. two papers have co-citation",
        "prompt node. two papers do not have co-citation",
    ]
    edge_text = [
        "feature edge. connected papers are cited together by other papers."
    ]
    prompt_node_text = [
        "prompt node. link prediction on the papers that are cited together"
    ]
    prompt_node_edge_text = [
        "prompt node. node classification on the paper's category"
    ]
    prompt_edge_text = ["prompt edge."]
    return (
        [new_data],
        [
            clean_text,
            label_text,
            edge_text,
            prompt_node_text,
            prompt_node_edge_text,
            prompt_edge_text,
            edge_label_text,
        ],
        None,
    )


class PubmedOFADataset(OFAPygDataset):
    def gen_data(self):
        cur_path = os.path.dirname(__file__)
        return get_data(os.path.join(cur_path, "pubmed.pt"))

    def add_text_emb(self, data_list, text_emb):
        data_list[0].x_text_feat = text_emb[0]
        data_list[0].label_text_feat = text_emb[1]
        data_list[0].edge_text_feat = text_emb[2]
        data_list[0].prompt_node_feat = text_emb[3]
        data_list[0].prompt_node_edge_feat = text_emb[4]
        data_list[0].prompt_edge_feat = text_emb[5]
        data_list[0].edge_label_feat = text_emb[6]
        return self.collate(data_list)

data/Pubmed/gen_data.py
LechengKong/OneForAll
import os.path as osp
import torch_geometric as pyg
import torch
from data.ofa_data import OFAPygDataset

entity2id = {}
entity_lst = []
text_lst = []
with open(osp.join(osp.dirname(__file__), "entity2text.txt"), "r") as f:
    lines = f.readlines()
    for line in lines:
        tmp = line.strip().split("\t")
        entity_lst.append(tmp[0])
        text_lst.append(tmp[1])

entity2id = {entity: i for i, entity in enumerate(entity_lst)}


def read_knowledge_graph(files):
    relation2id = {}

    converted_triplets = {}
    rel_list = []
    rel = len(relation2id)

    for file_type, file_path in files.items():

        edges = []
        edge_types = []
        with open(file_path) as f:
            file_data = [line.split() for line in f.read().split("\n")[:-1]]

        for triplet in file_data:
            if triplet[1] not in relation2id:
                relation2id[triplet[1]] = rel
                rel_list.append(triplet[1])
                rel += 1

            edges.append(
                [
                    entity2id[triplet[0]],
                    entity2id[triplet[2]],
                ]
            )
            edge_types.append(relation2id[triplet[1]])
        converted_triplets[file_type] = [edges, edge_types]

    new_data = pyg.data.data.Data(
        x=torch.zeros([len(text_lst), 1]),
        edge_index=torch.tensor(converted_triplets["train"][0]).T,
        edge_types=torch.tensor(converted_triplets["train"][1]),
    )

    node_text = [
        "feature node. entity and entity description: " + ent
        for ent in text_lst
    ]
    edge_text = [
        "feature edge. relation between two entities. " + relation
        for relation in rel_list
    ] + [
        "feature edge. relation between two entities. the inverse relation of "
        + relation
        for relation in rel_list
    ]

    prompt_edge_text = ["prompt edge"]
    prompt_node_text = [
        "prompt node. relation type prediction between the connected entities."
    ]
    label_text = [
        "prompt node. relation between two entities. " + relation
        for relation in rel_list
    ]

    return (
        [new_data],
        [
            node_text,
            edge_text,
            label_text,
            prompt_edge_text,
            prompt_node_text,
        ],
        [converted_triplets, rel_list],
    )


class WN18RROFADataset(OFAPygDataset):
    def gen_data(self):
        cur_path = osp.dirname(__file__)
        names = ["train", "valid", "test"]
        name_dict = {n: osp.join(cur_path, n + ".txt") for n in names}
        return read_knowledge_graph(name_dict)

    def add_text_emb(self, data_list, text_emb):
        data_list[0].x_text_feat = text_emb[0]
        data_list[0].edge_text_feat = text_emb[1]
        data_list[0].edge_label_feat = text_emb[2]
        data_list[0].prompt_edge_feat = text_emb[3]
        data_list[0].prompt_text_feat = text_emb[4]
        return self.collate(data_list)

    def get_idx_split(self):
        return self.side_data[0]

data/WN18RR/gen_data.py
LechengKong/OneForAll
import os
import pandas as pd
from data.ofa_data import OFAPygDataset
from ogb.nodeproppred import PygNodePropPredDataset


def get_node_feature(path):
    # Node feature process
    nodeidx2paperid = pd.read_csv(
        os.path.join(path, "nodeidx2paperid.csv.gz"), index_col="node idx"
    )
    titleabs_url = (
        "https://snap.stanford.edu/ogb/data/misc/ogbn_arxiv/titleabs.tsv"
    )
    titleabs = pd.read_csv(
        titleabs_url,
        sep="\t",
        names=["paper id", "title", "abstract"],
        index_col="paper id",
    )

    titleabs = nodeidx2paperid.join(titleabs, on="paper id")
    text = (
        "feature node. paper title and abstract: "
        + titleabs["title"]
        + ". "
        + titleabs["abstract"]
    )
    node_text_lst = text.values

    return node_text_lst


def get_taxonomy(path):

    # read categories and description file
    f = open(os.path.join(path, "arxiv_CS_categories.txt"), "r").readlines()

    state = 0
    result = {"id": [], "name": [], "description": []}

    for line in f:
        if state == 0:
            assert line.strip().startswith("cs.")
            category = (
                "arxiv "
                + " ".join(line.strip().split(" ")[0].split(".")).lower()
            )  # e. g. cs lo
            name = line.strip()[7:-1]  # e. g. Logic in CS
            result["id"].append(category)
            result["name"].append(name)
            state = 1
            continue
        elif state == 1:
            description = line.strip()
            result["description"].append(description)
            state = 2
            continue
        elif state == 2:
            state = 0
            continue

    arxiv_cs_taxonomy = pd.DataFrame(result)

    return arxiv_cs_taxonomy


def get_label_feature(path):

    arxiv_cs_taxonomy = get_taxonomy(path)
    mapping_file = os.path.join(path, "labelidx2arxivcategeory.csv.gz")
    arxiv_categ_vals = pd.merge(
        pd.read_csv(mapping_file),
        arxiv_cs_taxonomy,
        left_on="arxiv category",
        right_on="id",
    )

    text = (
        "prompt node. literature category and description: "
        + arxiv_categ_vals["name"]
        + ". "
        + arxiv_categ_vals["description"]
    )
    label_text_lst = text.values

    return label_text_lst


class ArxivOFADataset(OFAPygDataset):
    def gen_data(self):
        pyg_data = PygNodePropPredDataset(
            name="ogbn-arxiv", root=self.data_dir
        )
        cur_path = os.path.dirname(__file__)
        node_texts = get_node_feature(cur_path).tolist()
        label_texts = get_label_feature(cur_path).tolist()
        edge_text = ["feature edge. citation"]
        prompt_text = [
            "prompt node. node classification of literature category"
        ]
        prompt_edge_text = ["prompt edge."]
        return (
            [pyg_data.data],
            [
                node_texts,
                label_texts,
                edge_text,
                prompt_text,
                prompt_edge_text,
            ],
            None,
        )

    def add_text_emb(self, data_list, text_emb):
        data_list[0].x_text_feat = text_emb[0]
        data_list[0].label_text_feat = text_emb[1]
        data_list[0].edge_text_feat = text_emb[2]
        data_list[0].prompt_text_feat = text_emb[3]
        data_list[0].prompt_edge_feat = text_emb[4]
        return self.collate(data_list)

data/arxiv/gen_data.py
LechengKong/OneForAll
import os
import torch
import torch_geometric as pyg
import numpy as np
import json
from data.ofa_data import OFAPygDataset

import pickle
from datasets import load_dataset
import pandas as pd
from data.chemblpre.gen_raw_graph import smiles2graph


def load_prompt_json():
    with open(
        os.path.join(os.path.dirname(__file__), "prompt_pretrain.json"), "rb"
    ) as f:
        prompt_text = json.load(f)
    return prompt_text["chembl"]


def get_label_texts(labels):
    label_texts = [None] * len(labels)
    for entry in labels:
        label_texts[labels[entry][0]] = (
            "prompt node. molecule property description. "
            + "The molecule is effective to the following assay. "
            + labels[entry][1][0][:-41]
        )
    return label_texts


def get_local_text():
    print("gen text")
    cache_dir = os.path.join(
        os.path.dirname(__file__), "../../cache_data/dataset"
    )
    data = load_dataset(
        "haitengzhao/molecule_property_instruction",
        cache_dir=cache_dir,
        split="chembl_pretraining",
    )
    data_dict = {
        "label": data["label"],
        "task_index": data["task_index"],
        "molecule_index": data["molecule_index"],
    }
    pd_data = pd.DataFrame.from_dict(data_dict)
    cls_data = pd_data[np.logical_not(pd.isna(pd_data["task_index"]))]
    cls_data["ori_index"] = np.arange(len(cls_data))

    group = cls_data.groupby("molecule_index")
    index = group.ori_index.first()
    tasks = group.task_index.agg(lambda x: x.str.cat(sep=","))
    labels = group.label.agg(lambda x: x.str.cat(sep=","))
    mol = [data[i]["graph"] for i in index]
    split = [data[i]["split"] for i in index]

    prompt_text = load_prompt_json()
    task2index = {k: [i, prompt_text[k]] for i, k in enumerate(prompt_text)}
    label_text = get_label_texts(task2index)
    graphs = []
    for i in range(len(mol)):
        graph = smiles2graph(mol[i])
        task_lst = [task2index[v][0] for v in tasks[i].split(",")]
        label_lst = [1 if v == "Yes" else 0 for v in labels[i].split(",")]
        cur_label = np.zeros(len(task2index))
        cur_label[:] = np.nan
        cur_label[task_lst] = label_lst
        graph["label"] = cur_label
        graph["split"] = split[i]
        graphs.append(graph)
    return graphs, label_text


def gen_graph():
    graphs, labels_features = get_local_text()
    print("gen graph")

    node_texts = []
    edge_texts = []
    data = []
    for g in graphs:
        node_texts += g["node_feat"]
        edge_texts += g["edge_feat"]
    unique_node_texts = set(node_texts)
    unique_edge_texts = set(edge_texts)
    u_node_texts_lst = list(unique_node_texts)
    u_edge_texts_lst = list(unique_edge_texts)
    node_texts2id = {v: i for i, v in enumerate(u_node_texts_lst)}
    edge_texts2id = {v: i for i, v in enumerate(u_edge_texts_lst)}
    split = {"train": [], "valid": [], "test": []}
    for i, g in enumerate(graphs):
        cur_nt_id = [node_texts2id[v] for v in g["node_feat"]]
        cur_et_id = [edge_texts2id[v] for v in g["edge_feat"]]
        data.append(
            pyg.data.data.Data(
                x=torch.tensor(cur_nt_id, dtype=torch.long),
                xe=torch.tensor(cur_et_id, dtype=torch.long),
                edge_index=torch.tensor(g["edge_list"], dtype=torch.long).T,
                y=torch.tensor(g["label"]),
            )
        )
        split[g["split"]].append(i)

    prompt_edge_text = [
        "prompt edge.",
        "prompt edge. edge for query graph that is our target",
        "prompt edge. edge for support graph that is an example",
    ]
    prompt_text = [
        "prompt node. graph classification on molecule property",
        "prompt node. few shot task node for graph classification that decides whether the query molecule belongs to the class of support molecules.",
    ]

    ret = (
        data,
        [
            u_node_texts_lst,
            u_edge_texts_lst,
            labels_features,
            prompt_edge_text,
            prompt_text,
        ],
        split,
    )
    return ret
    # element = mol.nodes(data="element")
    # node_attr = {i: chem_dict[v] + " atom" for i, v in element}
    # nx.set_node_attributes(mol, node_attr, "node_desc")
    # m_order.update([b[2]["order"] for b in mol.edges(data=True)])


class CHEMBLPREOFADataset(OFAPygDataset):
    def gen_data(self):
        pyg_graph, texts, split = gen_graph()
        return [d for d in pyg_graph], texts, split

    def add_text_emb(self, data_list, text_emb):
        data, slices = self.collate(data_list)
        data.node_embs = text_emb[0]
        data.edge_embs = text_emb[1]
        data.label_text_feat = text_emb[2]
        data.prompt_edge_feat = text_emb[3]
        data.prompt_text_feat = text_emb[4]
        return data, slices

    def get(self, index):
        data = super().get(index)
        node_feat = self.node_embs[data.x]
        edge_feat = self.edge_embs[data.xe]
        data.x_text_feat = node_feat
        data.xe_text_feat = edge_feat
        data.y = data.y.view(1, -1)
        return data

    def get_idx_split(self):
        return self.side_data


if __name__ == "__main__":
    g, label = gen_graph()
    print(g[0])

data/chemblpre/gen_data.py
LechengKong/OneForAll
import os
import numpy as np

from rdkit import Chem
from rdkit.Chem import rdmolops
import networkx as nx
import pickle
import pandas as pd
import torch

file1 = open(os.path.join(os.path.dirname(__file__), "id2element.csv"), "r")
Lines = file1.readlines()
chem_dict = {}
for line in Lines:
    line_split = line.strip().split("\t")
    chem_dict[int(line_split[0])] = line_split[2]

allowable_features_map = {
    "possible_atomic_num_dict": chem_dict,
    "possible_chirality_dict": {
        "CHI_UNSPECIFIED": "unspecified",
        "CHI_TETRAHEDRAL_CW": "tetrahedral clockwise",
        "CHI_TETRAHEDRAL_CCW": "tetrahedral counter-clockwise",
        "CHI_OTHER": "other",
        "misc": "misc",
    },
    "possible_degree_list": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, "misc"],
    "possible_formal_charge_list": [
        -5,
        -4,
        -3,
        -2,
        -1,
        0,
        1,
        2,
        3,
        4,
        5,
        "misc",
    ],
    "possible_numH_list": [0, 1, 2, 3, 4, 5, 6, 7, 8, "misc"],
    "possible_number_radical_e_list": [0, 1, 2, 3, 4, "misc"],
    "possible_hybridization_list": [
        "SP",
        "SP2",
        "SP3",
        "SP3D",
        "SP3D2",
        "misc",
    ],
    "possible_is_aromatic_list": [False, True],
    "possible_is_in_ring_list": [False, True],
    "possible_bond_type_list": [
        "SINGLE",
        "DOUBLE",
        "TRIPLE",
        "AROMATIC",
        "misc",
    ],
    "possible_bond_stereo_dict": {
        "STEREONONE": "none",
        "STEREOZ": "Z",
        "STEREOE": "E",
        "STEREOCIS": "CIS",
        "STEREOTRANS": "TRANS",
        "STEREOANY": "ANY",
    },
    "possible_is_conjugated_list": [False, True],
}


def ReorderCanonicalRankAtoms(mol):
    order = tuple(
        zip(
            *sorted(
                [(j, i) for i, j in enumerate(Chem.CanonicalRankAtoms(mol))]
            )
        )
    )[1]
    mol_renum = Chem.RenumberAtoms(mol, order)
    return mol_renum, order


def get_chem_id2name():
    file1 = open("./id2element.csv", "r")
    Lines = file1.readlines()
    chem_dict = {}
    for line in Lines:
        line_split = line.strip().split(",")
        chem_dict[line_split[0]] = line_split[2]
    return chem_dict


def atom_to_feature(atom):
    """
    Converts rdkit atom object to feature list of indices
    :param mol: rdkit atom object
    :return: list
    """
    atom_feature = [
        chem_dict[int(atom.GetAtomicNum())],
        "atomic number is " + str(atom.GetAtomicNum()),
        allowable_features_map["possible_chirality_dict"][
            str(atom.GetChiralTag())
        ]
        + " chirality",
        "degree of " + str(atom.GetTotalDegree()),
        "formal charge of " + str(atom.GetFormalCharge()),
        "num of hydrogen is " + str(atom.GetTotalNumHs()),
        "num of radical electrons is " + str(atom.GetNumRadicalElectrons()),
        "hybridization is " + str(atom.GetHybridization()),
        "is aromatic" if atom.GetIsAromatic() else "not aromatric",
        "is in ring" if atom.IsInRing() else "not in ring",
    ]
    return "feature node. atom: " + " , ".join(atom_feature)


def bond_to_feature(bond):
    """
    Converts rdkit bond object to feature list of indices
    :param mol: rdkit bond object
    :return: list
    """
    bond_feature = [
        str(bond.GetBondType()) + " bond",
        "bond stereo is "
        + allowable_features_map["possible_bond_stereo_dict"][
            str(bond.GetStereo())
        ],
        "is conjugated" if bond.GetIsConjugated() else "not conjugated",
    ]
    return "feature edge. chemical bond. " + " , ".join(bond_feature)


def compute_cycle(mol):
    cycle_list = nx.cycle_basis(nx.Graph(rdmolops.GetAdjacencyMatrix(mol)))
    if len(cycle_list) == 0:
        cycle_length = 0
    else:
        cycle_length = max([len(j) for j in cycle_list])
    if cycle_length <= 6:
        cycle_length = 0
    else:
        cycle_length = cycle_length - 6
    cycle_score = -cycle_length
    return cycle_score


def smiles2graph(smiles_string, removeHs=True, reorder_atoms=False):
    """
    Converts SMILES string to graph Data object
    :input: SMILES string (str)
    :return: graph object
    """

    mol = Chem.MolFromSmiles(smiles_string)
    cycle_score = compute_cycle(mol)
    mol = mol if removeHs else Chem.AddHs(mol)
    if reorder_atoms:
        mol, _ = ReorderCanonicalRankAtoms(mol)

    # atoms
    atom_features_list = []
    for atom in mol.GetAtoms():
        atom_features_list.append(atom_to_feature(atom))

    # bonds
    edges_list = []
    edge_features_list = []
    bonds = mol.GetBonds()
    if len(bonds) == 0:
        edge_list = np.zeros((0, 2))
    else:
        for bond in mol.GetBonds():
            i = bond.GetBeginAtomIdx()
            j = bond.GetEndAtomIdx()

            edge_feature = bond_to_feature(bond)

            # add edges in both directions
            edges_list.append((i, j))
            edge_features_list.append(edge_feature)
            edges_list.append((j, i))
            edge_features_list.append(edge_feature)
        edge_list = np.array(edges_list)

    graph = dict()
    graph["edge_list"] = edge_list
    graph["edge_feat"] = edge_features_list
    graph["node_feat"] = atom_features_list
    graph["cycle"] = cycle_score

    return graph


def get_raw_graphs(data_path):
    arr = torch.load(
        os.path.join(os.path.dirname(__file__), "chembl_pretrain.pth"),
    )
    graphs = []
    for i, entry in enumerate(arr[0]):
        graph = smiles2graph(entry)
        graph["label"] = arr[1][i]
        graphs.append(graph)
    if not os.path.exists(data_path):
        os.mkdir(data_path)
    with open(os.path.join(data_path, "raw_graph.pkl"), "wb") as f:
        pickle.dump(graphs, f)

data/chemblpre/gen_raw_graph.py
LechengKong/OneForAll
import os
import torch
import torch_geometric as pyg
import numpy as np
import json
from data.ofa_data import OFAPygDataset

import pickle
from datasets import load_dataset
import pandas as pd
from data.chemblpre.gen_raw_graph import smiles2graph


def load_prompt_json():
    with open(
        os.path.join(os.path.dirname(__file__), "mol_label_desc.json"), "rb"
    ) as f:
        prompt_text = json.load(f)
    return prompt_text["hiv"]


def get_label_texts(labels):
    label_texts = [None] * int(len(labels) * 2)
    for entry in labels:
        label_texts[labels[entry][0] + len(labels)] = (
            "prompt node. molecule property description. "
            + "The molecule is effective to the following assay. "
            + labels[entry][1][0][:-41]
        )
        label_texts[labels[entry][0]] = (
            "prompt node. molecule property description. "
            + "The molecule is not effective to the following assay. "
            + labels[entry][1][0][:-41]
        )
    return label_texts


def get_local_text():
    print("gen text")
    cache_dir = os.path.join(
        os.path.dirname(__file__), "../../cache_data/dataset"
    )
    data = load_dataset(
        "haitengzhao/molecule_property_instruction",
        cache_dir=cache_dir,
        split="hiv",
    )
    data_dict = {
        "label": data["label"],
        "task_index": data["task_index"],
        "molecule_index": data["molecule_index"],
    }
    pd_data = pd.DataFrame.from_dict(data_dict)
    cls_data = pd_data[np.logical_not(pd.isna(pd_data["task_index"]))]
    cls_data["ori_index"] = np.arange(len(cls_data))

    group = cls_data.groupby("molecule_index")
    index = group.ori_index.first()
    tasks = group.task_index.agg(lambda x: x.str.cat(sep=","))
    labels = group.label.agg(lambda x: x.str.cat(sep=","))
    mol = [data[i]["graph"] for i in index]
    split = [data[i]["split"] for i in index]

    prompt_text = load_prompt_json()
    task2index = {k: [i, prompt_text[k]] for i, k in enumerate(prompt_text)}
    label_text = get_label_texts(task2index)
    graphs = []
    for i in range(len(mol)):
        graph = smiles2graph(mol[i])
        task_lst = [task2index[v][0] for v in tasks[i].split(",")]
        label_lst = [1 if v == "Yes" else 0 for v in labels[i].split(",")]
        cur_label = np.zeros(len(task2index))
        cur_label[:] = np.nan
        cur_label[task_lst] = label_lst
        graph["label"] = cur_label
        graph["split"] = split[i]
        graphs.append(graph)
    return graphs, label_text


def gen_graph():
    graphs, labels_features = get_local_text()
    print("gen graph")

    node_texts = []
    edge_texts = []
    data = []
    for g in graphs:
        node_texts += g["node_feat"]
        edge_texts += g["edge_feat"]
    unique_node_texts = set(node_texts)
    unique_edge_texts = set(edge_texts)
    u_node_texts_lst = list(unique_node_texts)
    u_edge_texts_lst = list(unique_edge_texts)
    node_texts2id = {v: i for i, v in enumerate(u_node_texts_lst)}
    edge_texts2id = {v: i for i, v in enumerate(u_edge_texts_lst)}
    split = {"train": [], "valid": [], "test": []}
    for i, g in enumerate(graphs):
        cur_nt_id = [node_texts2id[v] for v in g["node_feat"]]
        cur_et_id = [edge_texts2id[v] for v in g["edge_feat"]]
        data.append(
            pyg.data.data.Data(
                x=torch.tensor(cur_nt_id, dtype=torch.long),
                xe=torch.tensor(cur_et_id, dtype=torch.long),
                edge_index=torch.tensor(g["edge_list"], dtype=torch.long).T,
                y=torch.tensor(g["label"]),
            )
        )
        split[g["split"]].append(i)

    prompt_edge_text = [
        "prompt edge.",
        "prompt edge. edge for query graph that is our target",
        "prompt edge. edge for support graph that is an example",
    ]
    prompt_text = [
        "prompt node. graph classification on molecule property",
        "prompt node. few shot task node for graph classification that decides whether the query molecule belongs to the class of support molecules.",
    ]

    ret = (
        data,
        [
            u_node_texts_lst,
            u_edge_texts_lst,
            labels_features,
            prompt_edge_text,
            prompt_text,
        ],
        split,
    )
    return ret
    # element = mol.nodes(data="element")
    # node_attr = {i: chem_dict[v] + " atom" for i, v in element}
    # nx.set_node_attributes(mol, node_attr, "node_desc")
    # m_order.update([b[2]["order"] for b in mol.edges(data=True)])


class CHEMHIVOFADataset(OFAPygDataset):
    def gen_data(self):
        pyg_graph, texts, split = gen_graph()
        return [d for d in pyg_graph], texts, split

    def add_text_emb(self, data_list, text_emb):
        data, slices = self.collate(data_list)
        data.node_embs = text_emb[0]
        data.edge_embs = text_emb[1]
        data.label_text_feat = text_emb[2]
        data.prompt_edge_feat = text_emb[3]
        data.prompt_text_feat = text_emb[4]
        return data, slices

    def get(self, index):
        data = super().get(index)
        node_feat = self.node_embs[data.x]
        edge_feat = self.edge_embs[data.xe]
        data.x_text_feat = node_feat
        data.xe_text_feat = edge_feat
        data.y = data.y.view(1, -1)
        return data

    def get_idx_split(self):
        return self.side_data


if __name__ == "__main__":
    g, label = gen_graph()
    print(g[0])

data/chemhiv/gen_data.py
LechengKong/OneForAll
import os
import numpy as np

from rdkit import Chem
from rdkit.Chem import rdmolops
import networkx as nx
import pickle
import pandas as pd
import torch

file1 = open(os.path.join(os.path.dirname(__file__), "id2element.csv"), "r")
Lines = file1.readlines()
chem_dict = {}
for line in Lines:
    line_split = line.strip().split("\t")
    chem_dict[int(line_split[0])] = line_split[2]

allowable_features_map = {
    "possible_atomic_num_dict": chem_dict,
    "possible_chirality_dict": {
        "CHI_UNSPECIFIED": "unspecified",
        "CHI_TETRAHEDRAL_CW": "tetrahedral clockwise",
        "CHI_TETRAHEDRAL_CCW": "tetrahedral counter-clockwise",
        "CHI_OTHER": "other",
        "misc": "misc",
    },
    "possible_degree_list": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, "misc"],
    "possible_formal_charge_list": [
        -5,
        -4,
        -3,
        -2,
        -1,
        0,
        1,
        2,
        3,
        4,
        5,
        "misc",
    ],
    "possible_numH_list": [0, 1, 2, 3, 4, 5, 6, 7, 8, "misc"],
    "possible_number_radical_e_list": [0, 1, 2, 3, 4, "misc"],
    "possible_hybridization_list": [
        "SP",
        "SP2",
        "SP3",
        "SP3D",
        "SP3D2",
        "misc",
    ],
    "possible_is_aromatic_list": [False, True],
    "possible_is_in_ring_list": [False, True],
    "possible_bond_type_list": [
        "SINGLE",
        "DOUBLE",
        "TRIPLE",
        "AROMATIC",
        "misc",
    ],
    "possible_bond_stereo_dict": {
        "STEREONONE": "none",
        "STEREOZ": "Z",
        "STEREOE": "E",
        "STEREOCIS": "CIS",
        "STEREOTRANS": "TRANS",
        "STEREOANY": "ANY",
    },
    "possible_is_conjugated_list": [False, True],
}


def ReorderCanonicalRankAtoms(mol):
    order = tuple(
        zip(
            *sorted(
                [(j, i) for i, j in enumerate(Chem.CanonicalRankAtoms(mol))]
            )
        )
    )[1]
    mol_renum = Chem.RenumberAtoms(mol, order)
    return mol_renum, order


def get_chem_id2name():
    file1 = open("./id2element.csv", "r")
    Lines = file1.readlines()
    chem_dict = {}
    for line in Lines:
        line_split = line.strip().split(",")
        chem_dict[line_split[0]] = line_split[2]
    return chem_dict


def atom_to_feature(atom):
    """
    Converts rdkit atom object to feature list of indices
    :param mol: rdkit atom object
    :return: list
    """
    atom_feature = [
        chem_dict[int(atom.GetAtomicNum())],
        "atomic number is " + str(atom.GetAtomicNum()),
        allowable_features_map["possible_chirality_dict"][
            str(atom.GetChiralTag())
        ]
        + " chirality",
        "degree of " + str(atom.GetTotalDegree()),
        "formal charge of " + str(atom.GetFormalCharge()),
        "num of hydrogen is " + str(atom.GetTotalNumHs()),
        "num of radical electrons is " + str(atom.GetNumRadicalElectrons()),
        "hybridization is " + str(atom.GetHybridization()),
        "is aromatic" if atom.GetIsAromatic() else "not aromatric",
        "is in ring" if atom.IsInRing() else "not in ring",
    ]
    return "feature node. atom: " + " , ".join(atom_feature)


def bond_to_feature(bond):
    """
    Converts rdkit bond object to feature list of indices
    :param mol: rdkit bond object
    :return: list
    """
    bond_feature = [
        str(bond.GetBondType()) + " bond",
        "bond stereo is "
        + allowable_features_map["possible_bond_stereo_dict"][
            str(bond.GetStereo())
        ],
        "is conjugated" if bond.GetIsConjugated() else "not conjugated",
    ]
    return "feature edge. chemical bond. " + " , ".join(bond_feature)


def compute_cycle(mol):
    cycle_list = nx.cycle_basis(nx.Graph(rdmolops.GetAdjacencyMatrix(mol)))
    if len(cycle_list) == 0:
        cycle_length = 0
    else:
        cycle_length = max([len(j) for j in cycle_list])
    if cycle_length <= 6:
        cycle_length = 0
    else:
        cycle_length = cycle_length - 6
    cycle_score = -cycle_length
    return cycle_score


def smiles2graph(smiles_string, removeHs=True, reorder_atoms=False):
    """
    Converts SMILES string to graph Data object
    :input: SMILES string (str)
    :return: graph object
    """

    mol = Chem.MolFromSmiles(smiles_string)
    cycle_score = compute_cycle(mol)
    mol = mol if removeHs else Chem.AddHs(mol)
    if reorder_atoms:
        mol, _ = ReorderCanonicalRankAtoms(mol)

    # atoms
    atom_features_list = []
    for atom in mol.GetAtoms():
        atom_features_list.append(atom_to_feature(atom))

    # bonds
    edges_list = []
    edge_features_list = []
    bonds = mol.GetBonds()
    if len(bonds) == 0:
        edge_list = np.zeros((0, 2))
    else:
        for bond in mol.GetBonds():
            i = bond.GetBeginAtomIdx()
            j = bond.GetEndAtomIdx()

            edge_feature = bond_to_feature(bond)

            # add edges in both directions
            edges_list.append((i, j))
            edge_features_list.append(edge_feature)
            edges_list.append((j, i))
            edge_features_list.append(edge_feature)
        edge_list = np.array(edges_list)

    graph = dict()
    graph["edge_list"] = edge_list
    graph["edge_feat"] = edge_features_list
    graph["node_feat"] = atom_features_list
    graph["cycle"] = cycle_score

    return graph


def get_raw_graphs(data_path):
    arr = torch.load(
        os.path.join(os.path.dirname(__file__), "chembl_pretrain.pth"),
    )
    graphs = []
    for i, entry in enumerate(arr[0]):
        graph = smiles2graph(entry)
        graph["label"] = arr[1][i]
        graphs.append(graph)
    if not os.path.exists(data_path):
        os.mkdir(data_path)
    with open(os.path.join(data_path, "raw_graph.pkl"), "wb") as f:
        pickle.dump(graphs, f)

data/chemhiv/gen_raw_graph.py
LechengKong/OneForAll
import os
import torch
import torch_geometric as pyg
import numpy as np
import json
from data.ofa_data import OFAPygDataset

import pickle
from datasets import load_dataset
import pandas as pd
from data.chemblpre.gen_raw_graph import smiles2graph


def load_prompt_json():
    with open(
        os.path.join(os.path.dirname(__file__), "mol_label_desc.json"), "rb"
    ) as f:
        prompt_text = json.load(f)
    return prompt_text["pcba"]


def get_label_texts(labels):
    label_texts = [None] * len(labels)
    for entry in labels:
        label_texts[labels[entry][0]] = (
            "prompt node. molecule property description. "
            + "The molecule is effective to the following assay. "
            + labels[entry][1][0][:-41]
        )
    return label_texts


def get_local_text():
    print("gen text")
    cache_dir = os.path.join(
        os.path.dirname(__file__), "../../cache_data/dataset"
    )
    data = load_dataset(
        "haitengzhao/molecule_property_instruction",
        cache_dir=cache_dir,
        split="pcba",
    )
    data_dict = {
        "label": data["label"],
        "task_index": data["task_index"],
        "molecule_index": data["molecule_index"],
    }
    pd_data = pd.DataFrame.from_dict(data_dict)
    cls_data = pd_data[np.logical_not(pd.isna(pd_data["task_index"]))]
    cls_data["ori_index"] = np.arange(len(cls_data))

    group = cls_data.groupby("molecule_index")
    index = group.ori_index.first()
    tasks = group.task_index.agg(lambda x: x.str.cat(sep=","))
    labels = group.label.agg(lambda x: x.str.cat(sep=","))
    mol = [data[i]["graph"] for i in index]
    split = [data[i]["split"] for i in index]

    prompt_text = load_prompt_json()
    task2index = {k: [i, prompt_text[k]] for i, k in enumerate(prompt_text)}
    label_text = get_label_texts(task2index)
    graphs = []
    for i in range(len(mol)):
        graph = smiles2graph(mol[i])
        task_lst = [task2index[v][0] for v in tasks[i].split(",")]
        label_lst = [1 if v == "Yes" else 0 for v in labels[i].split(",")]
        cur_label = np.zeros(len(task2index))
        cur_label[:] = np.nan
        cur_label[task_lst] = label_lst
        graph["label"] = cur_label
        graph["split"] = split[i]
        graphs.append(graph)
    return graphs, label_text


def gen_graph():
    graphs, labels_features = get_local_text()
    print("gen graph")

    node_texts = []
    edge_texts = []
    data = []
    for g in graphs:
        node_texts += g["node_feat"]
        edge_texts += g["edge_feat"]
    unique_node_texts = set(node_texts)
    unique_edge_texts = set(edge_texts)
    u_node_texts_lst = list(unique_node_texts)
    u_edge_texts_lst = list(unique_edge_texts)
    node_texts2id = {v: i for i, v in enumerate(u_node_texts_lst)}
    edge_texts2id = {v: i for i, v in enumerate(u_edge_texts_lst)}
    split = {"train": [], "valid": [], "test": []}
    for i, g in enumerate(graphs):
        cur_nt_id = [node_texts2id[v] for v in g["node_feat"]]
        cur_et_id = [edge_texts2id[v] for v in g["edge_feat"]]
        data.append(
            pyg.data.data.Data(
                x=torch.tensor(cur_nt_id, dtype=torch.long),
                xe=torch.tensor(cur_et_id, dtype=torch.long),
                edge_index=torch.tensor(g["edge_list"], dtype=torch.long).T,
                y=torch.tensor(g["label"]),
            )
        )
        split[g["split"]].append(i)

    prompt_edge_text = [
        "prompt edge.",
        "prompt edge. edge for query graph that is our target",
        "prompt edge. edge for support graph that is an example",
    ]
    prompt_text = [
        "prompt node. graph classification on molecule property",
        "prompt node. few shot task node for graph classification that decides whether the query molecule belongs to the class of support molecules.",
    ]

    ret = (
        data,
        [
            u_node_texts_lst,
            u_edge_texts_lst,
            labels_features,
            prompt_edge_text,
            prompt_text,
        ],
        split,
    )
    return ret
    # element = mol.nodes(data="element")
    # node_attr = {i: chem_dict[v] + " atom" for i, v in element}
    # nx.set_node_attributes(mol, node_attr, "node_desc")
    # m_order.update([b[2]["order"] for b in mol.edges(data=True)])


class CHEMPCBAOFADataset(OFAPygDataset):
    def gen_data(self):
        pyg_graph, texts, split = gen_graph()
        return [d for d in pyg_graph], texts, split

    def add_text_emb(self, data_list, text_emb):
        data, slices = self.collate(data_list)
        data.node_embs = text_emb[0]
        data.edge_embs = text_emb[1]
        data.label_text_feat = text_emb[2]
        data.prompt_edge_feat = text_emb[3]
        data.prompt_text_feat = text_emb[4]
        return data, slices

    def get(self, index):
        data = super().get(index)
        node_feat = self.node_embs[data.x]
        edge_feat = self.edge_embs[data.xe]
        data.x_text_feat = node_feat
        data.xe_text_feat = edge_feat
        data.y = data.y.view(1, -1)
        return data

    def get_idx_split(self):
        return self.side_data


if __name__ == "__main__":
    g, label = gen_graph()
    print(g[0])

data/chempcba/gen_data.py
LechengKong/OneForAll
import os
import numpy as np

from rdkit import Chem
from rdkit.Chem import rdmolops
import networkx as nx
import pickle
import pandas as pd
import torch

file1 = open(os.path.join(os.path.dirname(__file__), "id2element.csv"), "r")
Lines = file1.readlines()
chem_dict = {}
for line in Lines:
    line_split = line.strip().split("\t")
    chem_dict[int(line_split[0])] = line_split[2]

allowable_features_map = {
    "possible_atomic_num_dict": chem_dict,
    "possible_chirality_dict": {
        "CHI_UNSPECIFIED": "unspecified",
        "CHI_TETRAHEDRAL_CW": "tetrahedral clockwise",
        "CHI_TETRAHEDRAL_CCW": "tetrahedral counter-clockwise",
        "CHI_OTHER": "other",
        "misc": "misc",
    },
    "possible_degree_list": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, "misc"],
    "possible_formal_charge_list": [
        -5,
        -4,
        -3,
        -2,
        -1,
        0,
        1,
        2,
        3,
        4,
        5,
        "misc",
    ],
    "possible_numH_list": [0, 1, 2, 3, 4, 5, 6, 7, 8, "misc"],
    "possible_number_radical_e_list": [0, 1, 2, 3, 4, "misc"],
    "possible_hybridization_list": [
        "SP",
        "SP2",
        "SP3",
        "SP3D",
        "SP3D2",
        "misc",
    ],
    "possible_is_aromatic_list": [False, True],
    "possible_is_in_ring_list": [False, True],
    "possible_bond_type_list": [
        "SINGLE",
        "DOUBLE",
        "TRIPLE",
        "AROMATIC",
        "misc",
    ],
    "possible_bond_stereo_dict": {
        "STEREONONE": "none",
        "STEREOZ": "Z",
        "STEREOE": "E",
        "STEREOCIS": "CIS",
        "STEREOTRANS": "TRANS",
        "STEREOANY": "ANY",
    },
    "possible_is_conjugated_list": [False, True],
}


def ReorderCanonicalRankAtoms(mol):
    order = tuple(
        zip(
            *sorted(
                [(j, i) for i, j in enumerate(Chem.CanonicalRankAtoms(mol))]
            )
        )
    )[1]
    mol_renum = Chem.RenumberAtoms(mol, order)
    return mol_renum, order


def get_chem_id2name():
    file1 = open("./id2element.csv", "r")
    Lines = file1.readlines()
    chem_dict = {}
    for line in Lines:
        line_split = line.strip().split(",")
        chem_dict[line_split[0]] = line_split[2]
    return chem_dict


def atom_to_feature(atom):
    """
    Converts rdkit atom object to feature list of indices
    :param mol: rdkit atom object
    :return: list
    """
    atom_feature = [
        chem_dict[int(atom.GetAtomicNum())],
        "atomic number is " + str(atom.GetAtomicNum()),
        allowable_features_map["possible_chirality_dict"][
            str(atom.GetChiralTag())
        ]
        + " chirality",
        "degree of " + str(atom.GetTotalDegree()),
        "formal charge of " + str(atom.GetFormalCharge()),
        "num of hydrogen is " + str(atom.GetTotalNumHs()),
        "num of radical electrons is " + str(atom.GetNumRadicalElectrons()),
        "hybridization is " + str(atom.GetHybridization()),
        "is aromatic" if atom.GetIsAromatic() else "not aromatric",
        "is in ring" if atom.IsInRing() else "not in ring",
    ]
    return "feature node. atom: " + " , ".join(atom_feature)


def bond_to_feature(bond):
    """
    Converts rdkit bond object to feature list of indices
    :param mol: rdkit bond object
    :return: list
    """
    bond_feature = [
        str(bond.GetBondType()) + " bond",
        "bond stereo is "
        + allowable_features_map["possible_bond_stereo_dict"][
            str(bond.GetStereo())
        ],
        "is conjugated" if bond.GetIsConjugated() else "not conjugated",
    ]
    return "feature edge. chemical bond. " + " , ".join(bond_feature)


def compute_cycle(mol):
    cycle_list = nx.cycle_basis(nx.Graph(rdmolops.GetAdjacencyMatrix(mol)))
    if len(cycle_list) == 0:
        cycle_length = 0
    else:
        cycle_length = max([len(j) for j in cycle_list])
    if cycle_length <= 6:
        cycle_length = 0
    else:
        cycle_length = cycle_length - 6
    cycle_score = -cycle_length
    return cycle_score


def smiles2graph(smiles_string, removeHs=True, reorder_atoms=False):
    """
    Converts SMILES string to graph Data object
    :input: SMILES string (str)
    :return: graph object
    """

    mol = Chem.MolFromSmiles(smiles_string)
    cycle_score = compute_cycle(mol)
    mol = mol if removeHs else Chem.AddHs(mol)
    if reorder_atoms:
        mol, _ = ReorderCanonicalRankAtoms(mol)

    # atoms
    atom_features_list = []
    for atom in mol.GetAtoms():
        atom_features_list.append(atom_to_feature(atom))

    # bonds
    edges_list = []
    edge_features_list = []
    bonds = mol.GetBonds()
    if len(bonds) == 0:
        edge_list = np.zeros((0, 2))
    else:
        for bond in mol.GetBonds():
            i = bond.GetBeginAtomIdx()
            j = bond.GetEndAtomIdx()

            edge_feature = bond_to_feature(bond)

            # add edges in both directions
            edges_list.append((i, j))
            edge_features_list.append(edge_feature)
            edges_list.append((j, i))
            edge_features_list.append(edge_feature)
        edge_list = np.array(edges_list)

    graph = dict()
    graph["edge_list"] = edge_list
    graph["edge_feat"] = edge_features_list
    graph["node_feat"] = atom_features_list
    graph["cycle"] = cycle_score

    return graph


def get_raw_graphs(data_path):
    arr = torch.load(
        os.path.join(os.path.dirname(__file__), "chembl_pretrain.pth"),
    )
    graphs = []
    for i, entry in enumerate(arr[0]):
        graph = smiles2graph(entry)
        graph["label"] = arr[1][i]
        graphs.append(graph)
    if not os.path.exists(data_path):
        os.mkdir(data_path)
    with open(os.path.join(data_path, "raw_graph.pkl"), "wb") as f:
        pickle.dump(graphs, f)

data/chempcba/gen_raw_graph.py
LechengKong/OneForAll
import json
import os.path
from abc import abstractmethod

# utils
import numpy as np
import torch
from ogb.nodeproppred import PygNodePropPredDataset
from torch.utils.data import BatchSampler
from torch.utils.data import DataLoader, Dataset
from torch_geometric.datasets import (
    Planetoid,
    Amazon,
    CoraFull,
    Coauthor,
    WikiCS,
)
from tqdm.autonotebook import trange

import torch_geometric.transforms as T
from torch_geometric.utils import to_undirected

from sentence_transformers import SentenceTransformer
from transformers import (
    LlamaForCausalLM,
    LlamaTokenizer,
    LlamaModel,
    LlamaConfig,
)


"""
For Debug:
from data.gen_dataset import load_dataset
    dataset = load_dataset('arxiv','llama2')
"""


def load_dataset(name, lm_name):
    """
    Return PyG dataset
    encoded node feature -> 'xdesc'
    encoded category feature -> 'ydesc'
    """
    dataset_dir = os.path.join(
        os.path.dirname(__file__), "../cache_data", name
    )
    if not os.path.exists(dataset_dir):
        os.mkdir(dataset_dir)
    available_datasets = ["arxiv", "wikics"]

    # Load dataset from PyG
    if name == "arxiv":
        from data.arxiv.gen_data import get_text

        dataset = PygNodePropPredDataset(name="ogbn-" + name, root=dataset_dir)
        # change arxiv to undirected graph
        edge_index = to_undirected(dataset._data.edge_index)
        set_dataset_attr(
            dataset, "edge_index", edge_index, edge_index.shape[1]
        )
    elif name == "wikics":
        from data.wikics.gen_data import get_text

        dataset = WikiCS(root=dataset_dir)
    else:
        raise ValueError(
            f"Unknown dataset: {name}. Please choose from {available_datasets}."
        )

    if not (
        os.path.exists(
            os.path.join(dataset_dir, f"{name}_{lm_name}_node_emb.pth")
        )
        or os.path.exists(
            os.path.join(dataset_dir, f"{name}_{lm_name}_class_emb.pth")
        )
    ):

        # Load raw text
        node_text_lst, class_text_lst = get_text(name)

        # Generate text embeddings
        print("Generating Embeddings.")
        node_embeddings, class_embeddings = text_model(
            node_text_lst, class_text_lst, name=lm_name
        )

        # Save embeddings
        torch.save(
            node_embeddings,
            os.path.join(dataset_dir, f"{name}_{lm_name}_node_emb.pth"),
        )
        torch.save(
            class_embeddings,
            os.path.join(dataset_dir, f"{name}_{lm_name}_class_emb.pth"),
        )

    else:
        print("Loading Embeddings.")
        node_embeddings = torch.load(
            os.path.join(dataset_dir, f"{name}_{lm_name}_node_emb.pth"),
            map_location="cpu",
        )
        class_embeddings = torch.load(
            os.path.join(dataset_dir, f"{name}_{lm_name}_class_emb.pth"),
            map_location="cpu",
        )

    # add node text feature embedding
    set_dataset_attr(
        dataset, "xdesc", node_embeddings, node_embeddings.shape[0]
    )
    # add label text feature embedding
    set_dataset_attr(
        dataset, "ydesc", class_embeddings, class_embeddings.shape[0]
    )

    return dataset


def text_model(node_text_lst, class_text_lst, name="ST"):
    if name == "ST":
        model = SentenceTransformer(
            "multi-qa-distilbert-cos-v1", device="cuda:0", cache_folder=os.path.join(os.path.dirname(__file__), "../cache_data")
        )
        node_embeddings = model.encode(
            node_text_lst,
            batch_size=128,
            show_progress_bar=True,
            convert_to_tensor=True,
        ).cpu()
        class_embeddings = model.encode(
            class_text_lst, show_progress_bar=True, convert_to_tensor=True
        ).cpu()
    elif name == "llama2":
        # model_path = os.path.join(os.path.abspath(os.path.dirname(os.getcwd())), 'llama-2-7b')
        model_path = os.path.join(os.getcwd(), "llama-2-7b")
        tokenizer = LlamaTokenizer.from_pretrained(model_path, device="cuda:0")
        model = LlamaModel.from_pretrained(model_path).to("cuda:0")

        node_embeddings = get_llama_emb(
            tokenizer, model, node_text_lst, batch_size=1
        )
        class_embeddings = get_llama_emb(
            tokenizer, model, class_text_lst, batch_size=1
        )
    else:
        raise ValueError(f"Unknown language model: {name}.")

    return node_embeddings, class_embeddings


def get_llama_emb(
    tokenizer,
    model,
    text_lst,
    batch_size,
    show_progress_bar=True,
    convert_to_tensor=True,
):

    # Add EOS token for padding
    tokenizer.pad_token = tokenizer.eos_token
    all_embeddings = []

    for start_index in trange(
        0,
        len(text_lst),
        batch_size,
        desc="Batches",
        disable=not show_progress_bar,
    ):
        sentences_batch = text_lst[start_index : start_index + batch_size]
        input_ids = tokenizer(
            sentences_batch, return_tensors="pt", padding=True
        ).input_ids.to("cuda:0")
        transformer_output = model(input_ids)

        # No gradients on word_embeddings
        word_embeddings = transformer_output[0].detach()
        sentence_embeddings = word_embeddings.mean(dim=1)
        all_embeddings.extend(sentence_embeddings)

    if convert_to_tensor:
        all_embeddings = torch.stack(all_embeddings)

    # input_ids = tokenizer(text_lst, return_tensors="pt", padding=True).input_ids
    #
    # transformer_output = model(input_ids)
    #
    # # No gradients on word_embeddings
    # word_embeddings = transformer_output[0].detach()
    # sentence_embeddings = word_embeddings.mean(dim=1)

    return all_embeddings


def set_dataset_attr(dataset, name, value, size):
    dataset._data_list = None
    dataset.data[name] = value
    if dataset.slices is not None:
        dataset.slices[name] = torch.tensor([0, size], dtype=torch.long)

data/gen_dataset.py
LechengKong/OneForAll
import os
import os.path as osp
import torch
import numpy as np
from tqdm import tqdm
from torch_geometric.data import InMemoryDataset, Dataset

from gp.utils.utils import SmartTimer


def safe_mkdir(path):
    if not osp.exists(path):
        os.mkdir(path)


def pth_safe_save(obj, path):
    if obj is not None:
        torch.save(obj, path)


def pth_safe_load(path):
    if osp.exists(path):
        return torch.load(path)
    return None


class OFAPygDataset(InMemoryDataset):
    def __init__(
        self,
        name,
        root="./cache_data",
        sentence_encoder=None,
        load_text=False,
        load_feat=True,
        transform=None,
        pre_transform=None,
        meta_dict=None,
    ):

        self.name = name
        self.root = root
        self.data_dir = osp.join(self.root, self.name)
        self.sentence_encoder = sentence_encoder
        safe_mkdir(self.data_dir)
        super().__init__(self.data_dir, transform, pre_transform)

        if load_text:
            self.texts = torch.load(self.processed_paths[1])

        self.data, self.slices = torch.load(self.processed_paths[0])
        self.side_data = pth_safe_load(self.processed_paths[2])

    def data2vec(self, data):
        if self.sentence_encoder is None:
            raise NotImplementedError("sentence encoder is not passed")
        if data is None:
            return None
        embeddings = self.sentence_encoder.encode(data).cpu()
        return embeddings

    @property
    def num_classes(self):
        return self.__num_classes__

    @property
    def raw_file_names(self):
        return []

    @property
    def processed_file_names(self):
        return ["geometric_data_processed.pt", "texts.pkl", "data.pt"]

    def text2feature(self, texts):
        if isinstance(texts[0], str):
            return self.data2vec(texts)
        return [self.text2feature(t) for t in texts]

    def gen_data(self):
        pass

    def add_text_emb(self, data_list, texts_emb):
        pass

    def process(self):
        data_list, texts, side_data = self.gen_data()

        texts_emb = self.text2feature(texts)
        torch.save(texts, self.processed_paths[1])
        if side_data is not None:
            torch.save(side_data, self.processed_paths[2])
        else:
            torch.save("No side data", self.processed_paths[2])
        data, slices = self.add_text_emb(data_list, texts_emb)

        print("Saving...")
        torch.save((data, slices), self.processed_paths[0])


class OFAPygSTDataset(Dataset):
    def __init__(
        self,
        name,
        root="./cache_data",
        sentence_encoder=None,
        load_text=False,
        load_feat=True,
        transform=None,
        pre_transform=None,
        meta_dict=None,
    ):

        self.name = name
        self.root = root
        self.data_dir = osp.join(self.root, self.name)
        self.sentence_encoder = sentence_encoder
        safe_mkdir(self.data_dir)
        super().__init__(self.data_dir, transform, pre_transform)

        if load_text:
            self.texts = torch.load(self.processed_paths[0])

        self.side_data = pth_safe_load(self.processed_paths[1])
        self.global_data = pth_safe_load(self.processed_paths[2])

        self.convert_data()

    def data2vec(self, data):
        if self.sentence_encoder is None:
            raise NotImplementedError("sentence encoder is not passed")
        if data is None:
            return None
        embeddings = self.sentence_encoder.encode(data).cpu()

        return embeddings

    @property
    def num_classes(self):
        return self.__num_classes__

    @property
    def raw_file_names(self):
        return []

    @property
    def processed_file_names(self):
        return [
            "texts.pkl",
            "data.pt",
            "global_data.pt",
            "node_feat.npy",
            "edge_feat.npy",
        ]

    def len(self):
        return 0

    def text2feature(self, texts):
        if isinstance(texts[0], str):
            return self.data2vec(texts)
        return [self.text2feature(t) for t in texts]

    def gen_data(self):
        pass

    def add_text_emb(self, data_list, texts_emb):
        pass

    def convert_data(self):
        pass

    def process(self):
        data_list, texts, side_data = self.gen_data()

        texts_emb = self.text2feature(texts)
        torch.save(texts, self.processed_paths[0])
        if side_data is not None:
            torch.save(side_data, self.processed_paths[1])
        else:
            torch.save("No side data", self.processed_paths[1])
        data, global_data = self.add_text_emb(data_list, texts_emb)
        if global_data is not None:
            torch.save(global_data, self.processed_paths[2])
        else:
            torch.save("No global data", self.processed_paths[2])

        print("Saving...")

        node_memmap = np.memmap(
            self.processed_paths[3],
            dtype="float32",
            mode="w+",
            shape=data[0].shape,
        )
        node_memmap[:] = data[0]
        node_memmap.flush()

        edge_memmap = np.memmap(
            self.processed_paths[4],
            dtype="float32",
            mode="w+",
            shape=data[1].shape,
        )
        edge_memmap[:] = data[1]
        edge_memmap.flush()

    def get(self, idx):
        data = torch.load(self.processed_paths[idx + 3])
        return data

data/ofa_data.py
LechengKong/OneForAll
import json
import os
import functools
import pickle
from data.ofa_data import OFAPygDataset
from torch_geometric.datasets import WikiCS

# For debug
#  from data.wikics.gen_data import get_text
#  a,b = get_text('wikics')


def get_text(path):
    """
    Returns: node_text_lst, label_text_lst
    Node text format: "wikipedia entry name: xxx. entry content: xxxxx"
    Label text format: "wikipedia entry category: xxx"
    """
    with open(os.path.join(path, "metadata.json")) as json_file:
        raw_data = json.load(json_file)

    node_info = raw_data["nodes"]
    label_info = raw_data["labels"]
    node_text_lst = []
    label_text_lst = []

    # Process Node Feature
    for node in node_info:
        node_feature = (
            (
                "feature node. wikipedia entry name: "
                + node["title"]
                + ". entry content: "
                + functools.reduce(lambda x, y: x + " " + y, node["tokens"])
            )
            .lower()
            .strip()
        )
        node_text_lst.append(node_feature)

    # Process Label Feature
    for label in label_info.values():
        label_feature = (
            ("prompt node. wikipedia entry category: " + label).lower().strip()
        )
        label_text_lst.append(label_feature)

    return node_text_lst, label_text_lst


class WikiCSOFADataset(OFAPygDataset):
    def gen_data(self):
        pyg_data = WikiCS(root=self.data_dir)
        cur_path = os.path.dirname(__file__)
        node_texts, label_texts = get_text(cur_path)
        edge_text = ["feature edge. wikipedia page link"]
        prompt_text = [
            "prompt node. node classification of wikipedia entry category"
        ]
        prompt_edge_text = ["prompt edge."]
        return (
            [pyg_data.data],
            [
                node_texts,
                label_texts,
                edge_text,
                prompt_text,
                prompt_edge_text,
            ],
            None,
        )

    def add_text_emb(self, data_list, text_emb):
        data_list[0].x_text_feat = text_emb[0]
        data_list[0].label_text_feat = text_emb[1]
        data_list[0].edge_text_feat = text_emb[2]
        data_list[0].prompt_text_feat = text_emb[3]
        data_list[0].prompt_edge_feat = text_emb[4]
        return self.collate(data_list)

data/wikics/gen_data.py
LechengKong/OneForAll
import json
import os.path
from abc import abstractmethod
from torch.utils.data import BatchSampler
from torch.utils.data import DataLoader, Dataset
import torch
import numpy as np


class DataManager:
    @abstractmethod
    def get_data_loader(self, mode):
        pass


class FewShotDataManager(DataManager):
    """
    Return dataloader for train/val/test node idx.
    Example:
        data_manager = FewShotDataManager(g, params)
        train_dataloader = data_manager.get_dataloader(0)
        val_dataloader = data_manager.get_dataloader(1)
        test_dataloader = data_manager.get_dataloader(2)
        next(iter(train_dataloader)).shape: (n_way, k_shot + q_query)
    """

    def __init__(
        self,
        data,
        n_way,
        k_shot,
        q_query,
        class_split_ratio=None,
        class_split_lst=None,
        num_workers=0,
    ):
        super(FewShotDataManager, self).__init__()
        data.y = data.y.squeeze()
        self.n_way = n_way
        self.num_workers = num_workers
        self.dataset = FewShotDataset(
            data,
            k_shot + q_query,
            class_split_ratio,
            num_workers=self.num_workers,
            class_split_lst=class_split_lst,
        )
        # self.split = self.dataset.split

    def get_data_loader(self, mode):
        # mode: 0->train, 1->val, 2->test
        class_list = self.dataset.__getclass__(mode)
        sampler = EpisodeBatchSampler(self.n_way, class_list)
        data_loader_params = dict(
            batch_sampler=sampler,
            num_workers=self.num_workers,
            pin_memory=False,
        )
        data_loader = DataLoader(self.dataset, **data_loader_params)
        return data_loader

    def get_dataset(self):
        return self.dataset


class FewShotDataset(Dataset):
    def __init__(
        self, data, batch_size, class_split_ratio, num_workers, class_split_lst
    ):
        self.data = data
        self.batch_size = batch_size
        self.class_split_ratio = class_split_ratio
        self.num_workers = num_workers
        self.class_split_lst = class_split_lst

        self.unique_label = torch.unique(self.data.y)

        self.cls_split_lst = self.class_split()
        self.cls_dataloader = self.create_subdataloader()
        # self.split = self.get_split_index()

    def class_split(self):
        """
        Split class for train/val/test in meta learning setting.
        Save as list: [[train_class_index], [val_class_index], [test_class_index], [all_class_index]]
        """
        # TODO: change data_path
        # cls_split_file = self.args.data_path + '/' + self.args.dataset + '_class_split.json'
        # cls_split_file = os.path.join(os.path.dirname(__file__), 'cache_data', )
        #
        # if os.path.isfile(cls_split_file) and False:
        #     # load list if exists
        #     with open(cls_split_file, 'rb') as f:
        #         cls_split_lst = json.load(f)
        #         print('Complete: Load class split info from %s .' % cls_split_file)
        if self.class_split_lst is not None:
            cls_split_lst = self.class_split_lst + [
                [cls for sublist in self.class_split_lst for cls in sublist]
            ]

        elif self.class_split_ratio is not None:
            # create list according to class_split_ratio and save
            label = self.unique_label.cpu().detach()
            valid_label_mask = torch.where(label >= 0)
            label = label[valid_label_mask]
            selected_labels = []
            if hasattr(self.data, "non_cs_labels"):
                print("Ignore overlapped labels in mag240m with ogbn-arxiv.")
                label = torch.tensor(self.data.non_cs_labels)
                selected_labels = self.data.cs_labels
            # if CoraFull dataset, ignore 68,69 label since they only have 15/29 samples
            # if label.size(0) == 70:
            #     label = label[:-2]
            # randomly shuffle
            label = label.index_select(0, torch.randperm(label.shape[0]))
            train_class, val_class, test_class = torch.split(
                label, self.class_split_ratio
            )
            cls_split_lst = [
                train_class.tolist(),
                val_class.tolist(),
                test_class.tolist(),
                label.tolist(),
                selected_labels,
            ]

        label = self.unique_label.cpu().detach()
        valid_label_mask = torch.where(label >= 0)
        label = label[valid_label_mask]
        # Just for unify with Prodigy (ogbn-arxiv setting)
        if label.shape[0] == 40:
            # TRAIN_LABELS = [0, 1, 2, 3, 4, 5, 7, 8, 9, 13, 15, 17, 18, 20, 22, 23, 24, 25, 26, 27, 28, 29, 31, 32, 33,
            #                 39]
            # VAL_LABELS = [6, 12, 16, 19, 30, 35, 38]
            # TEST_LABELS = [10, 11, 14, 21, 34, 36, 37]

            TRAIN_LABELS = [
                6,
                21,
                22,
                3,
                36,
                11,
                28,
                2,
                7,
                19,
                34,
                29,
                9,
                8,
                25,
                31,
                1,
                23,
                38,
                32,
            ]
            VAL_LABELS = [35, 24, 0, 27, 12, 13, 26, 20, 14, 4]
            TEST_LABELS = [17, 39, 10, 5, 16, 15, 18, 37, 30, 33]
            cls_split_lst = [
                TRAIN_LABELS,
                VAL_LABELS,
                TEST_LABELS,
                TRAIN_LABELS + VAL_LABELS + TEST_LABELS,
            ]

        if label.shape[0] == 237:
            TRAIN_LABELS = [
                44,
                109,
                9,
                144,
                106,
                73,
                172,
                52,
                211,
                113,
                226,
                191,
                29,
                148,
                111,
                98,
                21,
                186,
                228,
                119,
                185,
                66,
                199,
                112,
                174,
                72,
                207,
                85,
                200,
                115,
                84,
                60,
                125,
                61,
                12,
                99,
                149,
                27,
                78,
                181,
                10,
                128,
                30,
                42,
                45,
                100,
                89,
                4,
                180,
                161,
                230,
                5,
                130,
                192,
                3,
                15,
                175,
                131,
                156,
                145,
                190,
                110,
                219,
                2,
                222,
                203,
                79,
                122,
                81,
                47,
                206,
                164,
                176,
                38,
                13,
                77,
                223,
                64,
                123,
                0,
                146,
                105,
                193,
                37,
                157,
                153,
                82,
                34,
                124,
                187,
                197,
                103,
                201,
                232,
                32,
                215,
                165,
                104,
                48,
                171,
                177,
                76,
                120,
                196,
                132,
                209,
                135,
                19,
                212,
                213,
                208,
                169,
                133,
                65,
                170,
                235,
                46,
                51,
                121,
                204,
                134,
                167,
                126,
                33,
                138,
                162,
                95,
                63,
                229,
                184,
                36,
                182,
                101,
                127,
                152,
                140,
                54,
                16,
                155,
                102,
                80,
                220,
            ]

            VAL_LABELS = [
                168,
                14,
                83,
                205,
                142,
                69,
                236,
                118,
                159,
                231,
                136,
                43,
                18,
                68,
                53,
                90,
                94,
                41,
                93,
                116,
                195,
                225,
                25,
                216,
                198,
                74,
                58,
                210,
                17,
                49,
                147,
                92,
                158,
                160,
                75,
                141,
                20,
                96,
                31,
                137,
                117,
                11,
                67,
                214,
                88,
                91,
                24,
            ]

            TEST_LABELS = [
                97,
                218,
                227,
                86,
                217,
                39,
                202,
                87,
                221,
                178,
                40,
                194,
                1,
                71,
                150,
                114,
                56,
                107,
                224,
                179,
                166,
                183,
                50,
                143,
                234,
                154,
                129,
                59,
                55,
                23,
                7,
                8,
                108,
                151,
                22,
                139,
                233,
                173,
                26,
                188,
                35,
                57,
                62,
                70,
                189,
                6,
                28,
                163,
            ]
            cls_split_lst = [
                TRAIN_LABELS,
                VAL_LABELS,
                TEST_LABELS,
                TRAIN_LABELS + VAL_LABELS + TEST_LABELS,
            ]

            # with open(cls_split_file, 'w') as f:
            #     json.dump(cls_split_lst, f)
            #     print('Complete: Save class split info to %s .' % cls_split_file)
        # if len(cls_split_lst[1]) > 0:
        # print(cls_split_lst[:])
        return cls_split_lst

    def label_to_index(self) -> (dict, torch.tensor):
        """
        Generate a dictionary mapping labels to index list
        :return: dictionary: {label: [list of index]}
        """
        label = self.unique_label
        label2index = {}
        remove_label_list = []
        for i in label:
            idx = torch.nonzero(self.data.y == i)
            if idx.shape[0] < self.batch_size * 2:
                remove_label_list.append(int(i))
                label2index[int(i)] = None
            else:
                label2index[int(i)] = idx.squeeze()
        if len(remove_label_list) > 0:
            print(f"Remove invalid labels {len(remove_label_list)}.")
        self.invalid_labels = set(remove_label_list)

        return label2index, label

    def create_subdataloader(self):
        """
        :return: list of subdataloaders for each class i
        """
        label2index, label = self.label_to_index()
        cls_dataloader = []
        cls_dataloader_params = dict(
            batch_size=self.batch_size,
            shuffle=True,
            num_workers=self.num_workers,
            pin_memory=False,
        )
        for c in label:
            idx = label2index[int(c)]
            if idx is None:
                cls_dataloader.append(False)
            else:
                cls_dataset = ClassDataset(label2index[int(c)])
                cls_dataloader.append(
                    DataLoader(cls_dataset, **cls_dataloader_params)
                )

        return cls_dataloader

    def get_split_index(self):
        """
        :return: dictionary that contains the node index for each split
        """
        label2index, label = self.label_to_index()
        cls_split_lst = self.cls_split_lst
        split = {"train": [], "valid": [], "test": []}

        exclude_labels = []
        for c in label:
            if c in cls_split_lst[0]:
                split["train"].extend(
                    [int(idx) for idx in label2index[int(c)]]
                )
            elif c in cls_split_lst[1]:
                split["valid"].extend(
                    [int(idx) for idx in label2index[int(c)]]
                )
            elif c in cls_split_lst[2]:
                split["test"].extend([int(idx) for idx in label2index[int(c)]])
            else:
                exclude_labels.append(c)

        # print("Ignore labels: " % exclude_labels)

        return split

    def __getitem__(self, class_index):
        node_id = next(iter(self.cls_dataloader[class_index]))
        return next(iter(self.cls_dataloader[class_index])), class_index

    def __len__(self):
        # mode = 0 -> train; 1 -> validation; 2 -> test
        return len(self.unique_label)

    def __getclass__(self, mode):
        # return available classes under current mode (train/val/test)
        # print(self.cls_split_lst)
        class_list = self.cls_split_lst[mode]
        valid_labels = [
            label for label in class_list if label not in self.invalid_labels
        ]
        return valid_labels


class EpisodeBatchSampler(object):
    def __init__(self, n_way, class_list):
        # TODO: change value of episode to some variables
        self.episode = 1
        self.n_way = n_way
        self.class_list = class_list

    def __len__(self):
        return self.episode

    def __iter__(self):
        for i in range(self.episode):
            batch_class = []
            # Don't change task_num for OFA
            task_num = 1
            for j in range(task_num):
                batch_class.append(
                    np.random.choice(
                        self.class_list, self.n_way, replace=False
                    )
                )
            yield np.concatenate(batch_class)


class ClassDataset(Dataset):
    def __init__(self, label_index):
        self.label_index = label_index

    def __getitem__(self, i):
        return self.label_index[i]

    def __len__(self):
        return self.label_index.shape[0]

fs_datamanager.py
LechengKong/OneForAll
"""The ``gp`` package includes datasets and GNN models for
efficient graph learning.
"""

__version__ = "0.0.0"

gp/__init__.py
LechengKong/OneForAll
from typing import Union, List, Any, Optional, Dict
from lightning.pytorch import LightningDataModule
from torch.utils.data import DataLoader, RandomSampler
from gp.utils.datasets import DatasetWithCollate
from torch_geometric.data import Dataset as PygDataset
from torch_geometric.loader import DataLoader as PygDataloader
from torch.utils.data import Dataset


class DataWithMeta:
    def __init__(
        self,
        data: Dataset,
        batch_size: int,
        state_name: Optional[str] = None,
        feat_dim: int = 0,
        metric: Optional[str] = None,
        classes: Union[int, List[int]] = 2,
        is_regression: bool = False,
        meta_data: Any = None,
        sample_size: Optional[int] = -1,
    ):
        self.data = data
        self.batch_size = batch_size
        self.state_name = state_name
        self.feat_dim = feat_dim
        self.meta_data = meta_data
        self.metric = metric
        self.sample_size = sample_size
        self.classes = classes
        if isinstance(classes, list):
            self.num_tasks = len(classes)
        else:
            self.num_tasks = None
        self.is_regression = is_regression

    def pred_dim(self):
        if self.is_regression:
            return 1
        if self.num_tasks is not None:
            return self.num_tasks
        return self.classes


class DataModule(LightningDataModule):
    def __init__(
        self,
        data: Dict[str, DataWithMeta],
        num_workers: int = 4,
    ):
        super().__init__()
        self.datasets = data
        self.num_workers = num_workers

    def create_dataloader(
        self,
        data: Dataset,
        size: int,
        batch_size: int,
        drop_last: bool = True,
        shuffle: bool = True,
        num_workers: int = 0,
    ):
        sampler = None
        if size > 0:
            sampler = RandomSampler(data, num_samples=size, replacement=True)
            shuffle = False
        if isinstance(data, DatasetWithCollate):
            return DataLoader(
                data,
                batch_size,
                shuffle,
                sampler,
                num_workers=num_workers,
                collate_fn=data.get_collate_fn(),
                drop_last=drop_last,
            )
        if isinstance(data, PygDataset):
            return PygDataloader(
                data,
                batch_size,
                shuffle,
                sampler=sampler,
                num_workers=num_workers,
                drop_last=drop_last,
            )

    def train_dataloader(self):
        return self.create_dataloader(
            self.datasets["train"].data,
            self.datasets["train"].sample_size,
            self.datasets["train"].batch_size,
            num_workers=self.num_workers,
        )

    def val_dataloader(self):
        if isinstance(self.datasets["val"], list):
            data_list = []
            for val_data in self.datasets["val"]:
                data_list.append(
                    self.create_dataloader(
                        val_data.data,
                        val_data.sample_size,
                        val_data.batch_size,
                        drop_last=False,
                        shuffle=False,
                        num_workers=self.num_workers,
                    )
                )
            return data_list
        else:
            return [
                self.create_dataloader(
                    self.datasets["val"].data,
                    self.datasets["val"].sample_size,
                    self.datasets["val"].batch_size,
                    drop_last=False,
                    shuffle=False,
                    num_workers=self.num_workers,
                )
            ]

    def test_dataloader(self):
        if isinstance(self.datasets["test"], list):
            data_list = []
            for test_data in self.datasets["test"]:
                data_list.append(
                    self.create_dataloader(
                        test_data.data,
                        test_data.sample_size,
                        test_data.batch_size,
                        drop_last=False,
                        shuffle=False,
                        num_workers=self.num_workers,
                    )
                )
            return data_list
        else:
            return [
                self.create_dataloader(
                    self.datasets["test"].data,
                    self.datasets["test"].sample_size,
                    self.datasets["test"].batch_size,
                    drop_last=False,
                    shuffle=False,
                    num_workers=self.num_workers,
                )
            ]

gp/lightning/data_template.py
LechengKong/OneForAll
import torch
import os.path as osp
import copy
from typing import Any, Callable, Optional, Literal, List, Union
from torchmetrics import MeanAbsoluteError, Accuracy, AUROC, MeanMetric, Metric
from torchmetrics.text import BLEUScore
from itertools import chain


def classification_func(func, output, batch):
    output = output.view(-1, batch.num_classes[0])
    return func(output, batch.y.view(-1).to(torch.long))


def similarity_func(func, output, batch):
    return -func(output[0], output[1], dim=-1).mean()


def flat_binary_func(func, output, batch):
    labels = batch.bin_labels[batch.true_nodes_mask]
    valid_ind = labels == labels
    return func(output.view(-1)[valid_ind], labels[valid_ind])

def flat_binary_func_fs(func, output, batch):
    labels = batch.bin_labels
    valid_ind = labels == labels
    return func(output.view(-1)[valid_ind], labels[valid_ind])


def binary_auc_func(func, output, batch):
    output = output.view(-1, batch.num_classes[0])
    # score = torch.sigmoid(output)[:, -1]
    score = torch.nn.functional.softmax(output, dim=-1)[:, -1]
    return func(score, batch.y[:, -1].view(-1))


def regression_func(func, output, batch):
    return func(output.view(-1), batch.y.view(-1))


def batch_contrast(inp1, inp2, temp):
    anchor_emb = inp1.repeat_interleave(len(inp2), dim=0)
    neg_emb = inp2.repeat([len(inp1), 1])
    sim = torch.sum(anchor_emb * neg_emb, dim=-1) / temp
    sim = torch.exp(sim).view(len(inp1), len(inp2))
    print(sim)
    return -torch.log(torch.diagonal(sim) / sim.sum(dim=-1)).mean()


def get_contrast_func(temprature, cl_mode):
    if cl_mode == "twoview":

        def contrast_loss_func(loss, output, batch):
            return batch_contrast(output[0], output[1], temprature)

    elif cl_mode == "self":

        def contrast_loss_func(loss, output, batch):
            return batch_contrast(output, output, temprature)

    return contrast_loss_func

def generation_func(loss, output, batch):
    y = batch.y.to(torch.long).view(-1)
    return loss(output, y)

def BLEU_func(metric, output, batch):
    combine_texts = batch.combine_texts
    combine_texts = list(chain.from_iterable(combine_texts))
    combine_texts = [[text] for text in combine_texts]
    output_texts = batch.output_texts
    output_texts = list(chain.from_iterable(output_texts))
    output_texts = [[text] for text in output_texts]
    print(combine_texts[0])
    print(output[0])
    print("=" * 89)
    print(combine_texts[1])
    print(output[1])
    print("=" * 89)
    print(combine_texts[2])
    print(output[2])
    print("=" * 89)

    return metric(output, output_texts)


class EvalKit(torch.nn.Module):
    def __init__(
        self,
        metric_name: Union[str, List[str]],
        evlter: Any,
        loss: Any,
        evlter_func: Union[Callable, List[Callable]] = None,
        loss_func: Callable = None,
        val_monitor_state: Optional[str] = "valid",
        test_monitor_state: Optional[str] = "test",
        eval_mode: Literal["min", "max"] = "min",
        exp_prefix: str = "",
        eval_state: List[str] = ["train_eval", "test", "valid"],
    ):
        super().__init__()
        self.eval_states = eval_state
        self.loss = loss
        self.eval_mode = eval_mode
        self.val_monitor_state = val_monitor_state
        self.test_monitor_state = test_monitor_state
        self.exp_prefix = exp_prefix
        self.evlters = torch.nn.ModuleDict()
        self.loss_func = loss_func
        self.evlter_func = {}
        self.metric_name = {}
        for i, state in enumerate(eval_state):
            if isinstance(evlter, Metric):
                self.metric_name[state] = osp.join(
                    exp_prefix, state, metric_name
                )
                self.evlters[state] = copy.deepcopy(evlter)
                self.evlter_func[state] = evlter_func
            else:
                self.metric_name[state] = osp.join(
                    exp_prefix, state, metric_name[i]
                )
                self.evlters[state] = evlter[i]
                self.evlter_func[state] = evlter_func[i]

        self.val_metric = self.metric_name[self.val_monitor_state]
        self.test_metric = self.metric_name[self.test_monitor_state]

    def compute_loss(self, output: Any, batch: Any):
        return self.loss_func(self.loss, output, batch)

    def has_eval_state(self, state: str):
        return state in self.eval_states

    def get_evlter(self, state: str):
        return self.evlters[state]

    def eval_step(self, output: Any, batch: Any, state: str):
        evlter = self.get_evlter(state)
        return self.evlter_func[state](evlter, output, batch)

    def eval_epoch(self, state: str):
        evlter = self.get_evlter(state)
        return evlter.compute()

    def eval_reset(self, state: str):
        evlter = self.get_evlter(state)
        evlter.reset()

    def get_metric_name(self, state: str):
        return self.metric_name[state]


def prepare_mae(exp_name, eval_state=["train_eval", "test", "valid"]):
    evlter = MeanAbsoluteError()
    loss = torch.nn.L1Loss()
    return EvalKit(
        "mae",
        evlter,
        loss,
        regression_func,
        regression_func,
        exp_prefix=exp_name,
        eval_state=eval_state,
    )


def prepare_auc(exp_name, eval_state=["train_eval", "test", "valid"]):
    evlter = AUROC(task="binary")
    loss = torch.nn.CrossEntropyLoss()
    return EvalKit(
        "auc",
        evlter,
        loss,
        binary_auc_func,
        classification_func,
        eval_mode="max",
        exp_prefix=exp_name,
        eval_state=eval_state,
    )


def prepare_acc(
    exp_name, eval_state=["train_eval", "test", "valid"], **kwargs
):
    loss = torch.nn.CrossEntropyLoss()
    evlter = Accuracy(task="multiclass", num_classes=kwargs["num_class"])

    return EvalKit(
        "acc",
        evlter,
        loss,
        classification_func,
        classification_func,
        eval_mode="max",
        exp_prefix=exp_name,
        eval_state=eval_state,
    )


def prepare_bin_acc(
    exp_name, eval_state=["train_eval", "test", "valid"], **kwargs
):
    loss = torch.nn.BCEWithLogitsLoss()
    evlter = Accuracy(task="multiclass", num_classes=kwargs["num_class"])

    return EvalKit(
        "acc",
        evlter,
        loss,
        classification_func,
        flat_binary_func,
        eval_mode="max",
        exp_prefix=exp_name,
        eval_state=eval_state,
        val_state="valid_0",
    )


def prepare_bin_auc(
    exp_name, eval_state=["train_eval", "test", "valid"], **kwargs
):
    loss = torch.nn.BCEWithLogitsLoss()
    evlter = AUROC(task="binary")

    return EvalKit(
        "auc",
        evlter,
        loss,
        binary_auc_func,
        flat_binary_func,
        eval_mode="max",
        exp_prefix=exp_name,
        eval_state=eval_state,
    )


def prepare_cl(exp_name, eval_state=["train_eval", "test", "valid"], **kwargs):
    evlter = MeanMetric()

    return EvalKit(
        "acc",
        evlter,
        None,
        get_contrast_func(kwargs["temprature"], kwargs["mode"]),
        get_contrast_func(kwargs["temprature"], kwargs["mode"]),
        eval_mode="max",
        exp_prefix=exp_name,
        eval_state=eval_state,
    )


class IdentityLoss(torch.nn.Module):
    def __init__(self):
        super(IdentityLoss, self).__init__()

    def forward(self, inputs, targets):
        return inputs


def prepare_generation(
    exp_name, eval_state=["test", "valid"], **kwargs
):
    loss = IdentityLoss()
    evlter = BLEUScore()
    return EvalKit(
        "BLEU",
        evlter,
        loss,
        BLEU_func,
        generation_func,
        eval_mode="min",
        exp_prefix=exp_name,
        eval_state=eval_state,
    )



available_metrics = {
    "acc": prepare_acc,
    "mae": prepare_mae,
    "auc": prepare_auc,
    "cl": prepare_cl,
    "generation": prepare_generation,
}


def prepare_metric(metric, params, exp_name, data, eval_state=None):
    if metric not in available_metrics:
        raise NotImplementedError(metric + " not a available metric")
    if eval_state is not None:
        return available_metrics[metric](params, exp_name, data, eval_state)
    else:
        return available_metrics[metric](params, exp_name, data)

gp/lightning/metric.py
LechengKong/OneForAll
import torch
import os.path as osp
from typing import Optional, Union, List
from lightning.pytorch import LightningModule
from gp.lightning.metric import EvalKit


class ExpConfig:
    def __init__(
        self,
        name,
        optimizer,
        opt_params=None,
        lr_scheduler=None,
        dataset_callback=None,
    ):
        self.name = name
        self.optimizer = optimizer
        self.opt_params = opt_params
        self.lr_scheduler = lr_scheduler
        self.train_state_name = ["train_eval"]
        self.val_state_name = ["valid"]
        self.test_state_name = ["test"]
        self.dataset_callback = dataset_callback

    @property
    def train_state_name(self):
        return self._train_state_name

    @property
    def val_state_name(self):
        return self._val_state_name

    @property
    def test_state_name(self):
        return self._test_state_name

    @train_state_name.setter
    def train_state_name(self, value: Union[str, List[str]]):
        if isinstance(value, str):
            self._train_state_name = [value]
        else:
            self._train_state_name = value

    @val_state_name.setter
    def val_state_name(self, value: Union[str, List[str]]):
        if isinstance(value, str):
            self._val_state_name = [value]
        else:
            self._val_state_name = value

    @test_state_name.setter
    def test_state_name(self, value: Union[str, List[str]]):
        if isinstance(value, str):
            self._test_state_name = [value]
        else:
            self._test_state_name = value

    def get_optimizer(self):
        return self.optimizer

    def get_scheduler(self):
        return self.lr_scheduler


class BaseTemplate(LightningModule):
    def __init__(
        self,
        exp_config: ExpConfig,
        model: torch.nn.Module,
        eval_kit: Optional[EvalKit] = None,
        name: str = "",
    ):

        super().__init__()

        self.exp_config = exp_config

        self.model = model
        self.name = name

        self.eval_kit = eval_kit

    def on_test_epoch_start(self):
        self.on_validation_epoch_start()

    def configure_optimizers(self):
        optimizer = self.exp_config.get_optimizer()
        optimizer_dict = {"optimizer": optimizer}
        if self.exp_config.get_scheduler() is not None:
            optimizer_dict["lr_scheduler"] = self.exp_config.get_scheduler()
        return optimizer_dict

    def compute_results(
        self, batch, batch_idx, step_name, log_loss=True, *args
    ):
        try:
            score = self(batch, *args)
            loss = self.eval_kit.compute_loss(score, batch)
        except RuntimeError as e:
            if "out of memory" in str(e):
                print("Ignoring OOM batch")
                loss = None
                score = None
            else:
                raise
        if loss is not None:
            self.log(
                osp.join(self.name, step_name, "loss"),
                loss,
                on_step=False,
                on_epoch=True,
                prog_bar=log_loss,
                batch_size=batch.batch_size
                if hasattr(batch, "batch_size")
                else len(batch),
            )
        with torch.no_grad():
            if self.eval_kit.has_eval_state(step_name):
                self.eval_kit.eval_step(score, batch, step_name)
        return score, loss

    def epoch_post_process(self, epoch_name):
        if self.eval_kit.has_eval_state(epoch_name):
            metric = self.eval_kit.eval_epoch(epoch_name)
            self.log(
                self.eval_kit.get_metric_name(epoch_name),
                metric,
                prog_bar=True,
            )
            self.eval_kit.eval_reset(epoch_name)
            return metric

    def training_step(self, batch, batch_idx, dataloader_idx=0):
        score, loss = self.compute_results(
            batch, batch_idx, self.exp_config.train_state_name[dataloader_idx]
        )
        return loss

    def on_train_epoch_end(self):
        for name in self.exp_config.train_state_name:
            self.epoch_post_process(name)

    def validation_step(self, batch, batch_idx, dataloader_idx=0):
        self.compute_results(
            batch,
            batch_idx,
            self.exp_config.val_state_name[dataloader_idx],
            log_loss=False,
        )

    def on_validation_epoch_end(self):
        cur_metric = []
        for name in self.exp_config.val_state_name:
            metric = self.epoch_post_process(name)
            if metric is not None:
                cur_metric.append(metric.cpu())
        if self.exp_config.dataset_callback is not None:
            self.exp_config.dataset_callback(cur_metric)

    def test_step(self, batch, batch_idx, dataloader_idx=0):
        self.compute_results(
            batch,
            batch_idx,
            self.exp_config.test_state_name[dataloader_idx],
            log_loss=False,
        )

    def on_test_epoch_end(self):
        for name in self.exp_config.test_state_name:
            self.epoch_post_process(name)

gp/lightning/module_template.py
LechengKong/OneForAll
import torch
import numpy as np
from lightning.pytorch import Trainer
from lightning.pytorch.callbacks.progress import TQDMProgressBar
from lightning.pytorch.callbacks import ModelCheckpoint

from gp.utils.utils import dict_res_summary
from gp.lightning.metric import EvalKit


def lightning_fit(
    logger,
    model,
    data_module,
    metrics: EvalKit,
    num_epochs,
    profiler=None,
    cktp_prefix="",
    load_best=True,
    prog_freq=20,
    test_rep=1,
    save_model=True,
    prog_bar=True,
    accelerator="auto",
    detect_anomaly=False,
    reload_freq=0,
    val_interval=1,
):
    callbacks = []
    if prog_bar:
        callbacks.append(TQDMProgressBar(refresh_rate=20))
    if save_model:
        callbacks.append(
            ModelCheckpoint(
                monitor=metrics.val_metric,
                mode=metrics.eval_mode,
                save_last=True,
                filename=cktp_prefix + "{epoch}-{step}",
            )
        )
    trainer = Trainer(
        accelerator=accelerator,
        devices=1 if torch.cuda.is_available() else 10,
        max_epochs=num_epochs,
        callbacks=callbacks,
        logger=logger,
        log_every_n_steps=prog_freq,
        profiler=profiler,
        enable_checkpointing=save_model,
        enable_progress_bar=prog_bar,
        detect_anomaly=detect_anomaly,
        reload_dataloaders_every_n_epochs=reload_freq,
        check_val_every_n_epoch=val_interval,
    )
    trainer.fit(model, datamodule=data_module)
    if load_best:
        model.load_state_dict(
            torch.load(trainer.checkpoint_callback.best_model_path)[
                "state_dict"
            ]
        )

    val_col = []
    for i in range(test_rep):
        val_col.append(
            trainer.validate(model, datamodule=data_module, verbose=False)[0]
        )

    val_res = dict_res_summary(val_col)
    for met in val_res:
        val_mean = np.mean(val_res[met])
        val_std = np.std(val_res[met])
        print("{}:{:f}±{:f}".format(met, val_mean, val_std))

    target_val_mean = np.mean(val_res[metrics.val_metric])
    target_val_std = np.std(val_res[metrics.val_metric])

    test_col = []
    for i in range(test_rep):
        test_col.append(
            trainer.test(model, datamodule=data_module, verbose=False)[0]
        )

    test_res = dict_res_summary(test_col)
    for met in test_res:
        test_mean = np.mean(test_res[met])
        test_std = np.std(test_res[met])
        print("{}:{:f}±{:f}".format(met, test_mean, test_std))

    target_test_mean = np.mean(test_res[metrics.test_metric])
    target_test_std = np.std(test_res[metrics.test_metric])
    return [target_val_mean, target_val_std], [
        target_test_mean,
        target_test_std,
    ]


def lightning_test(
    logger,
    model,
    data_module,
    metrics: EvalKit,
    model_dir: str,
    profiler=None,
    prog_freq=20,
    test_rep=1,
    prog_bar=True,
    accelerator="auto",
    detect_anomaly=False,
):
    callbacks = []
    if prog_bar:
        callbacks.append(TQDMProgressBar(refresh_rate=20))
    trainer = Trainer(
        accelerator=accelerator,
        devices=1 if torch.cuda.is_available() else 10,
        callbacks=callbacks,
        logger=logger,
        log_every_n_steps=prog_freq,
        profiler=profiler,
        enable_progress_bar=prog_bar,
        detect_anomaly=detect_anomaly,
    )
    # trainer.fit(model, datamodule=data_module)
    model.load_state_dict(torch.load(model_dir)["state_dict"])

    val_col = []
    for i in range(test_rep):
        val_col.append(
            trainer.validate(model, datamodule=data_module, verbose=False)[0]
        )

    val_res = dict_res_summary(val_col)
    for met in val_res:
        val_mean = np.mean(val_res[met])
        val_std = np.std(val_res[met])
        print("{}:{:f}±{:f}".format(met, val_mean, val_std))

    target_val_mean = np.mean(val_res[metrics.val_metric])
    target_val_std = np.std(val_res[metrics.val_metric])

    test_col = []
    for i in range(test_rep):
        test_col.append(
            trainer.test(model, datamodule=data_module, verbose=False)[0]
        )

    test_res = dict_res_summary(test_col)
    for met in test_res:
        test_mean = np.mean(test_res[met])
        test_std = np.std(test_res[met])
        print("{}:{:f}±{:f}".format(met, test_mean, test_std))

    target_test_mean = np.mean(test_res[metrics.test_metric])
    target_test_std = np.std(test_res[metrics.test_metric])
    return [target_val_mean, target_val_std], [
        target_test_mean,
        target_test_std,
    ]

gp/lightning/training.py
LechengKong/OneForAll
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn import Parameter
from torch_geometric.nn import MessagePassing
from torch_geometric.nn.inits import glorot, zeros
from torch_geometric.typing import Adj, OptTensor
from torch_geometric.utils import softmax, add_self_loops


def masked_edge_index(edge_index, edge_mask):
    if isinstance(edge_index, torch.Tensor):
        return edge_index[:, edge_mask]


class RGCNEdgeConv(MessagePassing):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        num_relations: int,
        aggr: str = "mean",
        **kwargs,
    ):
        kwargs.setdefault("aggr", aggr)
        super().__init__(**kwargs)  # "Add" aggregation (Step 5).
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.num_relations = num_relations

        self.weight = Parameter(
            torch.empty(self.num_relations, in_channels, out_channels)
        )

        self.root = Parameter(torch.empty(in_channels, out_channels))
        self.bias = Parameter(torch.empty(out_channels))

        self.reset_parameters()

    def reset_parameters(self):
        super().reset_parameters()
        glorot(self.weight)
        glorot(self.root)
        zeros(self.bias)

    def forward(
        self,
        x: OptTensor,
        xe: OptTensor,
        edge_index: Adj,
        edge_type: OptTensor = None,
    ):
        out = torch.zeros(x.size(0), self.out_channels, device=x.device)
        for i in range(self.num_relations):
            edge_mask = edge_type == i
            tmp = masked_edge_index(edge_index, edge_mask)

            h = self.propagate(tmp, x=x, xe=xe[edge_mask])
            out += h @ self.weight[i]

        out += x @ self.root
        out += self.bias

        return out

    def message(self, x_j, xe):
        # x_j has shape [E, out_channels]

        # Step 4: Normalize node features.
        return (x_j + xe).relu()


class RGATEdgeConv(RGCNEdgeConv):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        num_relations: int,
        aggr: str = "sum",
        heads=8,
        add_self_loops=False,
        share_att=False,
        **kwargs,
    ):
        self.heads = heads
        self.add_self_loops = add_self_loops
        self.share_att = share_att
        super().__init__(
            in_channels,
            out_channels,
            num_relations,
            aggr,
            node_dim=0,
            **kwargs,
        )
        self.lin_edge = nn.Linear(self.in_channels, self.out_channels)
        assert self.in_channels % heads == 0
        self.d_model = self.in_channels // heads
        if self.share_att:
            self.att = Parameter(torch.empty(1, heads, self.d_model))
        else:
            self.att = Parameter(
                torch.empty(self.num_relations, heads, self.d_model)
            )

        glorot(self.att)

    def forward(
        self,
        x: OptTensor,
        xe: OptTensor,
        edge_index: Adj,
        edge_type: OptTensor = None,
    ):
        out = torch.zeros((x.size(0), self.out_channels), device=x.device)

        if self.add_self_loops:
            num_nodes = x.size(0)
            edge_index, xe = add_self_loops(
                edge_index, xe, fill_value="mean", num_nodes=num_nodes
            )

        x_ = x.view(-1, self.heads, self.d_model)
        xe_ = self.lin_edge(xe).view(-1, self.heads, self.d_model)

        for i in range(self.num_relations):
            edge_mask = edge_type == i
            if self.add_self_loops:
                edge_mask = torch.cat(
                    [
                        edge_mask,
                        torch.ones(num_nodes, device=edge_mask.device).bool(),
                    ]
                )

            tmp = masked_edge_index(edge_index, edge_mask)

            h = self.propagate(tmp, x=x_, xe=xe_[edge_mask], rel_index=i)
            h = h.view(-1, self.in_channels)
            out += h @ self.weight[i]

        out += x @ self.root
        out += self.bias

        return out

    def message(self, x_j, xe, rel_index, index, ptr, size_i):
        # x_j has shape [E, out_channels]
        x = F.leaky_relu(x_j + xe)
        if self.share_att:
            att = self.att

        else:
            att = self.att[rel_index : rel_index + 1]

        alpha = (x * att).sum(dim=-1)
        alpha = softmax(alpha, index, ptr, size_i)

        return (x_j + xe) * alpha.unsqueeze(-1)

gp/nn/layer/pyg.py
LechengKong/OneForAll
"""Customized loss functions
"""

import torch


class InfoNCEloss(torch.nn.Module):
    def __init__(self) -> None:
        super().__init__()

    def forward(self, MI_MAT):
        n = len(MI_MAT)
        e_neg_mat = (
            MI_MAT.view(-1)[1:].view(n - 1, n + 1)[:, :-1].reshape(n, n - 1)
        )
        e_pos = torch.diagonal(MI_MAT)
        loss = -torch.mean(
            torch.log(torch.exp(e_pos) / torch.exp(e_neg_mat).sum(dim=-1))
        )
        return loss


class CCALoss(torch.nn.Module):
    def __init__(self, outdim_size=20):
        super().__init__()
        self.outdim_size = outdim_size

    def forward(self, H1, H2):
        """
        It is the loss function of CCA as introduced in the original paper.
        There can be other formulations.
        """

        r1 = 1e-5
        r2 = 1e-5
        eps = 1e-7

        H1, H2 = H1.t(), H2.t()
        # assert torch.isnan(H1).sum().item() == 0
        # assert torch.isnan(H2).sum().item() == 0

        o1 = o2 = H1.size(0)

        m = H1.size(1)
        #         print(H1.size())

        H1bar = H1 - H1.mean(dim=1).unsqueeze(dim=1)
        H2bar = H2 - H2.mean(dim=1).unsqueeze(dim=1)
        # assert torch.isnan(H1bar).sum().item() == 0
        # assert torch.isnan(H2bar).sum().item() == 0

        SigmaHat12 = (1.0 / (m - 1)) * torch.matmul(H1bar, H2bar.t())
        SigmaHat11 = (1.0 / (m - 1)) * torch.matmul(
            H1bar, H1bar.t()
        ) + r1 * torch.eye(o1, device=H1.device)
        SigmaHat22 = (1.0 / (m - 1)) * torch.matmul(
            H2bar, H2bar.t()
        ) + r2 * torch.eye(o2, device=H1.device)
        # print(SigmaHat11.sum())
        # print(SigmaHat22.sum())
        # print(SigmaHat12.sum())
        assert torch.isnan(SigmaHat11).sum().item() == 0
        assert torch.isnan(SigmaHat12).sum().item() == 0
        assert torch.isnan(SigmaHat22).sum().item() == 0

        [D1, V1] = torch.linalg.eigh(SigmaHat11)
        [D2, V2] = torch.linalg.eigh(SigmaHat22)
        v1, c1 = torch.unique(D1, return_counts=True)
        v2, c2 = torch.unique(D2, return_counts=True)
        if len(v1[c1 > 1]) > 0 or len(v2[c2 > 1]) > 0:
            print("sht")
            return (
                torch.tensor(0, requires_grad=True, dtype=torch.float),
                torch.eye(o1),
                torch.eye(o1),
            )
        assert torch.isnan(D1).sum().item() == 0
        assert torch.isnan(D2).sum().item() == 0
        assert torch.isnan(V1).sum().item() == 0
        assert torch.isnan(V2).sum().item() == 0

        # Added to increase stability
        posInd1 = torch.gt(D1, eps).nonzero()[:, 0]
        D1 = D1[posInd1]
        V1 = V1[:, posInd1]
        posInd2 = torch.gt(D2, eps).nonzero()[:, 0]
        D2 = D2[posInd2]
        V2 = V2[:, posInd2]
        # print(posInd1.size())
        # print(posInd2.size())

        SigmaHat11RootInv = torch.matmul(
            torch.matmul(V1, torch.diag(D1**-0.5)), V1.t()
        )
        SigmaHat22RootInv = torch.matmul(
            torch.matmul(V2, torch.diag(D2**-0.5)), V2.t()
        )

        Tval = torch.matmul(
            torch.matmul(SigmaHat11RootInv, SigmaHat12), SigmaHat22RootInv
        )
        #         print(Tval.size())

        # just the top self.outdim_size singular values are used
        trace_TT = torch.matmul(Tval.t(), Tval)
        trace_TT = torch.add(
            trace_TT, (torch.eye(trace_TT.shape[0]) * r1).to(H1.device)
        )  # regularization for more stability
        U, V = torch.linalg.eigh(trace_TT)
        U = torch.where(U > eps, U, (torch.ones(U.shape) * eps).to(H1.device))
        U = U.topk(self.outdim_size)[0]
        corr = torch.sum(torch.sqrt(U))
        U, S, V = torch.svd(Tval)
        U = U[:, : self.outdim_size]
        V = V[:, : self.outdim_size]
        U = torch.matmul(SigmaHat11RootInv, U)
        V = torch.matmul(SigmaHat22RootInv, V)
        return corr, U, V


class IDLoss(torch.nn.Module):
    def __init__(self) -> None:
        super().__init__()

    def forward(self, res):
        return res


class NegLogLoss(torch.nn.Module):
    def __init__(self, num_neg_samples) -> None:
        super().__init__()
        self.neg_sample = num_neg_samples

    def forward(self, res):
        score_mat = res.view(-1, self.neg_sample + 1)
        score_mat = torch.sigmoid(score_mat)
        loss = torch.mean(
            -torch.log(score_mat[:, 0])
            - (torch.log(1 - score_mat[:, 1:])).mean(dim=-1)
        )
        return loss


class FirstPosNegLoss(torch.nn.Module):
    def __init__(self, num_neg_samples) -> None:
        super().__init__()
        self.neg_sample = num_neg_samples
        self.loss = torch.nn.BCEWithLogitsLoss()

    def forward(self, res):
        score_mat = res.view(-1, self.neg_sample + 1)
        target = torch.zeros_like(score_mat)
        target[:, 0] = 1
        loss = self.loss(score_mat.flatten(), target.flatten())
        return loss


class MRRLoss(torch.nn.Module):
    def __init__(self, num_neg_samples) -> None:
        super().__init__()
        self.loss = torch.nn.MarginRankingLoss(15, reduction="sum")
        self.num_neg_samples = num_neg_samples

    def forward(self, res):
        scores_mat = res.view(-1, self.num_neg_samples + 1)
        score_pos = (
            scores_mat[:, 0]
            .unsqueeze(1)
            .repeat_interleave(self.num_neg_samples, dim=-1)
        )
        score_neg = scores_mat[:, 1:]
        loss = self.loss(
            score_pos,
            score_neg,
            torch.ones_like(score_pos).to(device=res.device),
        )
        return loss

gp/nn/loss.py
LechengKong/OneForAll
"""Base message-passing GNNs
"""
import torch
import torch.nn as nn
import torch.nn.functional as F

from gp.nn.models.util_model import MLP

from abc import ABCMeta, abstractmethod

from gp.utils.utils import SmartTimer

from torch_scatter import scatter


class MultiLayerMessagePassing(nn.Module, metaclass=ABCMeta):
    """Message passing GNN"""

    def __init__(
        self,
        num_layers,
        inp_dim,
        out_dim,
        drop_ratio=None,
        JK="last",
        batch_norm=True,
    ):
        """

        :param num_layers: layer number of GNN
        :type num_layers: int
        :param inp_dim: input feature dimension
        :type inp_dim: int
        :param out_dim: output dimension
        :type out_dim: int
        :param drop_ratio: layer-wise node dropout ratio, defaults to None
        :type drop_ratio: float, optional
        :param JK: jumping knowledge, should either be ["last","sum"],
        defaults to "last"
        :type JK: str, optional
        :param batch_norm: Use node embedding batch normalization, defaults
        to True
        :type batch_norm: bool, optional
        """
        super().__init__()
        self.num_layers = num_layers
        self.JK = JK
        self.inp_dim = inp_dim
        self.out_dim = out_dim

        self.drop_ratio = drop_ratio

        self.conv = torch.nn.ModuleList()

        if batch_norm:
            self.batch_norm = torch.nn.ModuleList()
            for layer in range(num_layers):
                self.batch_norm.append(torch.nn.BatchNorm1d(out_dim))
        else:
            self.batch_norm = None

        self.timer = SmartTimer(False)

    def build_layers(self):
        for layer in range(self.num_layers):
            if layer == 0:
                self.conv.append(self.build_input_layer())
            else:
                self.conv.append(self.build_hidden_layer())

    @abstractmethod
    def build_input_layer(self):
        pass

    @abstractmethod
    def build_hidden_layer(self):
        pass

    @abstractmethod
    def layer_forward(self, layer, message):
        pass

    @abstractmethod
    def build_message_from_input(self, g):
        pass

    @abstractmethod
    def build_message_from_output(self, g, output):
        pass

    def forward(self, g):
        h_list = []

        message = self.build_message_from_input(g)

        for layer in range(self.num_layers):
            # print(layer, h)
            h = self.layer_forward(layer, message)
            if self.batch_norm:
                h = self.batch_norm[layer](h)
            if layer != self.num_layers - 1:
                h = F.relu(h)
            if self.drop_ratio is not None:
                h = F.dropout(h, p=self.drop_ratio, training=self.training)
            message = self.build_message_from_output(g, h)
            h_list.append(h)

        if self.JK == "last":
            repr = h_list[-1]
        elif self.JK == "sum":
            repr = 0
            for layer in range(self.num_layers):
                repr += h_list[layer]
        else:
            repr = h_list
        return repr


class MultiLayerMessagePassingVN(MultiLayerMessagePassing):
    def __init__(
        self,
        num_layers,
        inp_dim,
        out_dim,
        drop_ratio=None,
        JK="last",
        batch_norm=True,
    ):
        super().__init__(
            num_layers, inp_dim, out_dim, drop_ratio, JK, batch_norm
        )

        self.virtualnode_embedding = torch.nn.Embedding(1, self.out_dim)
        torch.nn.init.constant_(self.virtualnode_embedding.weight.data, 0)

        self.virtualnode_mlp_list = torch.nn.ModuleList()
        for layer in range(self.num_layers - 1):
            self.virtualnode_mlp_list.append(
                MLP([self.out_dim, 2 * self.out_dim, self.out_dim])
            )

    def forward(self, g):
        h_list = []

        message = self.build_message_from_input(g)

        vnode_embed = self.virtualnode_embedding(
            torch.zeros(g.batch_size, dtype=torch.int).to(g.device)
        )

        batch_node_segment = torch.arange(
            g.batch_size, dtype=torch.long, device=g.device
        ).repeat_interleave(g.batch_num_nodes())

        for layer in range(self.num_layers):
            # print(layer, h)
            h = self.layer_forward(layer, message)
            if self.batch_norm:
                h = self.batch_norm[layer](h)
            if layer != self.num_layers - 1:
                h = F.relu(h)
            if self.drop_ratio is not None:
                h = F.dropout(h, p=self.drop_ratio, training=self.training)
            message = self.build_message_from_output(g, h)
            h_list.append(h)

            if layer < self.num_layers - 1:
                vnode_emb_temp = (
                    scatter(
                        h, batch_node_segment, dim=0, dim_size=g.batch_size
                    )
                    + vnode_embed
                )

                vnode_embed = F.dropout(
                    self.virtualnode_mlp_list[layer](vnode_emb_temp),
                    self.drop_ratio,
                    training=self.training,
                )

        if self.JK == "last":
            repr = h_list[-1]
        elif self.JK == "sum":
            repr = 0
            for layer in range(self.num_layers):
                repr += h_list[layer]
        elif self.JK == "cat":
            repr = torch.cat([h_list], dim=-1)
        return repr

gp/nn/models/GNN.py
LechengKong/OneForAll
"""gp extracts gnn logics from packages and allows customized
layers, which is particularly useful for fast prototyping.
gp also implements models in dgl and pytorch-geometrics.
gp is task-oriented with task predictors.
"""

gp/nn/models/__init__.py
LechengKong/OneForAll
import torch
from gp.nn.models.GNN import MultiLayerMessagePassing

from gp.nn.models.util_model import MLP
from dgl.nn.pytorch.conv import GINConv, RelGraphConv


class DGLGIN(MultiLayerMessagePassing):
    def __init__(
        self,
        num_layers,
        inp_dim,
        out_dim,
        drop_ratio=0,
        JK="last",
        batch_norm=True,
    ):
        super().__init__(
            num_layers, inp_dim, out_dim, drop_ratio, JK, batch_norm
        )
        self.build_layers()

    def build_input_layer(self):
        return GINConv(
            MLP(
                [self.inp_dim, 2 * self.inp_dim, self.out_dim],
                batch_norm=self.batch_norm is not None,
            ),
            learn_eps=True,
        )

    def build_hidden_layer(self):
        return GINConv(
            MLP(
                [self.out_dim, 2 * self.out_dim, self.out_dim],
                batch_norm=self.batch_norm is not None,
            ),
            learn_eps=True,
        )

    def build_message_from_input(self, g, input_feat="feat"):
        if isinstance(input_feat, str):
            h = g.ndata[input_feat]
        elif torch.is_tensor(input_feat):
            h = input_feat
        else:
            raise NotImplementedError("Not supported input type")
        return {"g": g, "h": h}

    def build_message_from_output(self, g, h):
        return {"g": g, "h": h}

    def layer_forward(self, layer, message):
        return self.conv[layer](message["g"], message["h"])


class DGLRGCN(MultiLayerMessagePassing):
    def __init__(
        self,
        num_layers,
        num_rels,
        inp_dim,
        out_dim,
        num_bases=4,
        drop_ratio=0,
        JK="last",
        batch_norm=True,
    ):
        super().__init__(
            num_layers, inp_dim, out_dim, drop_ratio, JK, batch_norm
        )
        self.num_rels = num_rels
        self.num_bases = num_bases
        self.build_layers()

    def build_input_layer(self):
        return RelGraphConv(
            self.inp_dim, self.out_dim, self.num_rels, num_bases=self.num_bases
        )

    def build_hidden_layer(self):
        return RelGraphConv(
            self.out_dim, self.out_dim, self.num_rels, num_bases=self.num_bases
        )

    def build_message_from_input(self, g):
        return {"g": g, "h": g.ndata["feat"], "e": g.edata["type"]}

    def build_message_from_output(self, g, h):
        return {"g": g, "h": h, "e": g.edata["type"]}

    def layer_forward(self, layer, message):
        return self.conv[layer](message["g"], message["h"], message["e"])

gp/nn/models/dgl.py
LechengKong/OneForAll
from gp.nn.models.GNN import MultiLayerMessagePassing

from gp.nn.models.util_model import MLP
from torch_geometric.nn.conv import GINConv, GINEConv, RGCNConv
import torch


class PyGGIN(MultiLayerMessagePassing):
    def __init__(
        self,
        num_layers,
        inp_dim,
        out_dim,
        drop_ratio=0,
        JK="last",
        batch_norm=True,
    ):
        super().__init__(
            num_layers, inp_dim, out_dim, drop_ratio, JK, batch_norm
        )
        self.build_layers()

    def build_input_layer(self):
        return GINConv(
            MLP(
                [self.inp_dim, 2 * self.inp_dim, self.out_dim],
                batch_norm=self.batch_norm is not None,
            ),
            train_eps=True,
        )

    def build_hidden_layer(self):
        return GINConv(
            MLP(
                [self.out_dim, 2 * self.out_dim, self.out_dim],
                batch_norm=self.batch_norm is not None,
            ),
            train_eps=True,
        )

    def build_message_from_input(self, g):
        return {"g": g.edge_index, "h": g.x}

    def build_message_from_output(self, g, h):
        return {"g": g.edge_index, "h": h}

    def layer_forward(self, layer, message):
        return self.conv[layer](message["h"], message["g"])


class PyGGINE(MultiLayerMessagePassing):
    def __init__(
        self,
        num_layers,
        inp_dim,
        out_dim,
        edge_dim,
        drop_ratio=0,
        JK="last",
        batch_norm=True,
    ):
        super().__init__(
            num_layers, inp_dim, out_dim, drop_ratio, JK, batch_norm
        )
        self.edge_dim = edge_dim
        self.build_layers()

    def build_input_layer(self):
        return GINEConv(
            MLP([self.inp_dim, self.inp_dim * 2, self.out_dim]),
            train_eps=True,
            edge_dim=self.edge_dim,
        )

    def build_hidden_layer(self):
        return GINEConv(
            MLP([self.out_dim, self.out_dim * 2, self.out_dim]),
            train_eps=True,
            edge_dim=self.edge_dim,
        )

    def build_message_from_input(self, g):
        return {"g": g.edge_index, "h": g.x, "e": g.edge_attr}

    def build_message_from_output(self, g, h):
        return {"g": g.edge_index, "h": g.x, "e": g.edge_attr}

    def layer_forward(self, layer, message):
        return self.conv[layer](message["h"], message["g"], message["e"])


class PyGRGCN(MultiLayerMessagePassing):
    def __init__(
        self,
        num_layers,
        num_rels,
        inp_dim,
        out_dim,
        num_bases=None,
        drop_ratio=0,
        JK="last",
        batch_norm=True,
    ):
        super().__init__(
            num_layers, inp_dim, out_dim, drop_ratio, JK, batch_norm
        )
        self.num_rels = num_rels
        self.num_bases = num_bases
        # self.layer_norms = torch.nn.ModuleList(
        #     [torch.nn.LayerNorm(out_dim) for _ in range(num_layers)])  # <-- Define LayerNorm modules
        self.build_layers()

    def build_input_layer(self):
        return RGCNConv(
            self.inp_dim, self.out_dim, self.num_rels, num_bases=self.num_bases
        )

    def build_hidden_layer(self):
        return RGCNConv(
            self.out_dim, self.out_dim, self.num_rels, num_bases=self.num_bases
        )

    def build_message_from_input(self, g):
        return {"g": g.edge_index, "h": g.x, "e": g.edge_type}

    def build_message_from_output(self, g, h):
        return {"g": g.edge_index, "h": h, "e": g.edge_type}

    def layer_forward(self, layer, message):
        # h = self.conv[layer](message["h"], message["g"], message["e"])
        # h = self.layer_norms[layer](h)
        # return h
        return self.conv[layer](message["h"], message["g"], message["e"])

gp/nn/models/pyg.py
LechengKong/OneForAll
"""Task specific predictor. Use GNN to generate node embedding
and produce task specific predictions.
"""

import torch
import torch.nn as nn
import dgl

from abc import ABCMeta, abstractmethod
from gp.nn.pooling import ReprIndexTransform


class BaseGNNEncoder(nn.Module, metaclass=ABCMeta):
    """Takes in a GNN as memeber variable. Implements forward
    and get_out_dim.
    """

    def __init__(self, gnn):
        super().__init__()

        self.gnn = gnn

    @abstractmethod
    def forward(self, g, pool_kwargs=None, gnn_kwargs=None):
        pass

    @abstractmethod
    def get_out_dim(self):
        pass


class BaseGraphEncoder(BaseGNNEncoder, metaclass=ABCMeta):
    def __init__(self, emb_dim, gnn):
        super().__init__(gnn)

        self.emb_dim = emb_dim

    def forward(self, g, pool_kwargs=None, gnn_kwargs=None):
        if gnn_kwargs is not None:
            repr = self.gnn(g, **gnn_kwargs)
        else:
            repr = self.gnn(g)

        if pool_kwargs is not None:
            return self.pool_from_graph(g, repr, **pool_kwargs)
        else:
            return self.pool_from_graph(g, repr)

    @abstractmethod
    def pool_from_graph(self, g, repr, **kwargs):
        pass


class GraphEncoder(BaseGraphEncoder):
    def __init__(self, emb_dim, gnn, pooling="sum"):
        super().__init__(emb_dim, gnn)
        self.pooling = "sum"

    def pool_from_graph(self, g, repr):
        repr = self.repr_post_process(g, repr, input)

        g.ndata["node_res"] = repr

        g_repr = dgl.readout_nodes(g, "node_res", op=self.pooling)
        return g_repr

    def repr_post_process(self, g, repr, input):
        return repr

    def get_out_dim(self):
        return self.gnn.out_dim


class BaseLinkEncoder(BaseGNNEncoder, metaclass=ABCMeta):
    def __init__(self, emb_dim, gnn):
        super().__init__(gnn)
        self.emb_dim = emb_dim
        self.use_only_embedding = False

    def embedding_only_mode(self, state=True):
        self.use_only_embedding = state

    def forward(self, g, head, tail, input, gnn_kwargs=None):
        if self.use_only_embedding:
            repr = g.ndata["repr"]
        else:
            repr = self.gnn(g, **gnn_kwargs)
        return self.pool_from_link(repr, head, tail, input)

    @abstractmethod
    def pool_from_link(self, repr, head, tail, input):
        pass


class LinkEncoder(BaseLinkEncoder):
    def __init__(self, emb_dim, gnn):
        super().__init__(emb_dim, gnn)
        self.node_repr_extractor = ReprIndexTransform(gnn.out_dim)

    def pool_from_link(self, repr, head, tail, input):
        repr_list = []
        repr_list.append(self.node_repr_extractor(repr, head))
        repr_list.append(self.node_repr_extractor(repr, tail))
        g_rep = torch.cat(repr_list, dim=1)
        return g_rep

    def get_out_dim(self):
        return self.gnn.out_dim * 2


class BaseNodeEncoder(BaseGNNEncoder):
    def __init__(self, emb_dim, gnn):
        super().__init__(gnn)
        self.emb_dim = emb_dim
        self.use_only_embedding = False

    def embedding_only_mode(self, state=True):
        self.use_only_embedding = state

    def forward(self, g, node, input, gnn_kwargs=None):
        if self.use_only_embedding:
            repr = g.ndata["repr"]
        else:
            repr = self.gnn(g, **gnn_kwargs)
        return self.pool_from_node(g, repr, node, input)

    @abstractmethod
    def pool_from_node(self, g, repr, node, input):
        pass


class NodeEncoder(BaseNodeEncoder):
    def __init__(self, emb_dim, gnn):
        super().__init__(emb_dim, gnn)
        self.node_repr_extractor = ReprIndexTransform(gnn.out_dim)

    def pool_from_node(self, g, repr, node, input):
        return self.node_repr_extractor(repr, node)

    def get_out_dim(self):
        return self.gnn.out_dim

gp/nn/models/task_predictor.py
LechengKong/OneForAll
from typing import Any, Callable, Dict, List, Optional, Union

import torch
import torch.nn.functional as F
from torch import Tensor
from gp.nn.resolver import activation_resolver


class MLP(torch.nn.Module):
    """
    MLP model modifed from pytorch geometric.
    """

    def __init__(
        self,
        channel_list: Optional[Union[List[int], int]] = None,
        dropout: Union[float, List[float]] = 0.0,
        act: Union[str, Callable, None] = "relu",
        act_first: bool = False,
        act_kwargs: Optional[Dict[str, Any]] = None,
        norm: bool = True,
        plain_last: bool = True,
        bias: Union[bool, List[bool]] = True,
        **kwargs,
    ):
        super().__init__()

        assert isinstance(channel_list, (tuple, list))
        assert len(channel_list) >= 2
        self.channel_list = channel_list

        self.act = activation_resolver(act, **(act_kwargs or {}))
        self.act_first = act_first
        self.plain_last = plain_last

        if isinstance(dropout, float):
            dropout = [dropout] * (len(channel_list) - 1)
            if plain_last:
                dropout[-1] = 0.0
        if len(dropout) != len(channel_list) - 1:
            raise ValueError(
                f"Number of dropout values provided ({len(dropout)} does not "
                f"match the number of layers specified "
                f"({len(channel_list)-1})"
            )
        self.dropout = dropout

        if isinstance(bias, bool):
            bias = [bias] * (len(channel_list) - 1)
        if len(bias) != len(channel_list) - 1:
            raise ValueError(
                f"Number of bias values provided ({len(bias)}) does not match "
                f"the number of layers specified ({len(channel_list)-1})"
            )

        self.lins = torch.nn.ModuleList()
        iterator = zip(channel_list[:-1], channel_list[1:], bias)
        for in_channels, out_channels, _bias in iterator:
            self.lins.append(
                torch.nn.Linear(in_channels, out_channels, bias=_bias)
            )

        self.norms = torch.nn.ModuleList()
        iterator = channel_list[1:-1] if plain_last else channel_list[1:]
        for hidden_channels in iterator:
            if norm is not None:
                norm_layer = torch.nn.BatchNorm1d(hidden_channels)
            else:
                norm_layer = torch.nn.Identity()
            self.norms.append(norm_layer)

        self.reset_parameters()

    @property
    def in_channels(self) -> int:
        r"""Size of each input sample."""
        return self.channel_list[0]

    @property
    def out_channels(self) -> int:
        r"""Size of each output sample."""
        return self.channel_list[-1]

    @property
    def num_layers(self) -> int:
        r"""The number of layers."""
        return len(self.channel_list) - 1

    def reset_parameters(self):
        r"""Resets all learnable parameters of the module."""
        for lin in self.lins:
            lin.reset_parameters()
        for norm in self.norms:
            if hasattr(norm, "reset_parameters"):
                norm.reset_parameters()

    def forward(self, x: Tensor) -> Tensor:
        r"""
        Args:
            x (torch.Tensor): The source tensor.
            return_emb (bool, optional): If set to :obj:`True`, will
                additionally return the embeddings before execution of to the
                final output layer. (default: :obj:`False`)
        """
        for i, (lin, norm) in enumerate(zip(self.lins, self.norms)):
            x = lin(x)
            if self.act is not None and self.act_first:
                x = self.act(x)
            x = norm(x)
            if self.act is not None and not self.act_first:
                x = self.act(x)
            x = F.dropout(x, p=self.dropout[i], training=self.training)

        if self.plain_last:
            x = self.lins[-1](x)
            x = F.dropout(x, p=self.dropout[-1], training=self.training)

        return x

    def __repr__(self) -> str:
        return f"{self.__class__.__name__}({str(self.channel_list)[1:-1]})"

gp/nn/models/util_model.py
LechengKong/OneForAll
"""Neural Network pooling and Transformaation functions
"""

import torch
import torch.nn as nn

from abc import ABCMeta, abstractmethod
from torch_scatter import scatter

from gp.nn.models.util_model import MLP
from gp.utils.utils import count_to_group_index


class Extractor(nn.Module, metaclass=ABCMeta):
    @abstractmethod
    def get_out_dim(self):
        pass

    @abstractmethod
    def forward(self):
        pass


class Pooler(Extractor):
    def __init__(self):
        super().__init__()

    def get_out_dim(self):
        return None


class Transform(Extractor):
    def __init__(self, emb_dim):
        super().__init__()
        self.emb_dim = emb_dim

    def get_out_dim(self):
        return self.emb_dim


class GDTransform(Transform):
    """compute node-level GD representation."""

    def __init__(self, emb_dim, gd_deg=True) -> None:
        super().__init__(emb_dim)
        self.gd_deg = gd_deg
        if gd_deg:
            self.mlp_combine_gd_deg = MLP([emb_dim + 1, 2 * emb_dim, emb_dim])
        self.mlp_combine_nei_gd = MLP([2 * emb_dim + 1, 4 * emb_dim, emb_dim])
        self.mlp_combine_node_nei = MLP([2 * emb_dim, 4 * emb_dim, emb_dim])

    def forward(
        self,
        repr,
        nodes,
        neighbors,
        neighbor_count,
        dist,
        gd,
        gd_count,
        gd_deg,
    ):
        """

        Arguments:
            repr {torch.tensor} -- N*d tensor of node representations
            nodes {torch.tensor} -- 1-d nodes of interests
            neighbors {torch.tensor} -- neighbors indices of nodes
            neighbor_count {torch.tensor} -- nodes[i]'s neighbor=
            neighbors[cumsum(neighbor_count[:i]):cumsum(neighbor_count[:i+1])]
            dist {torch.tensor} -- dist[i]=distance between neighbors[i] and
            its corresponding source in nodes.
            gd {torch.tensor} -- Vertical Geodesics indices of neighbors
            gd_count {torch.tensor} -- neighbors[i]'s vertical geodesics=
                    gd[cumsum(gd_count[i]):cumsum(gd_count[i+1])]
            gd_deg {torch.tensor} -- same shape as gd, vertical gd degrees.

        Returns:
            Vertical geodesics of each node in nodes.
        """
        neighbors_repr = repr[neighbors]
        gd_repr = repr[gd]
        if self.gd_deg:
            combined_gd_repr = self.mlp_combine_gd_deg(
                torch.cat([gd_repr, gd_deg.view(-1, 1)], dim=-1)
            )
        else:
            combined_gd_repr = gd_repr
        combined_gd_repr = scatter(
            combined_gd_repr,
            count_to_group_index(gd_count),
            dim=0,
            dim_size=len(gd_count),
        )
        combined_repr = self.mlp_combine_nei_gd(
            torch.cat(
                [combined_gd_repr, neighbors_repr, dist.view(-1, 1)], dim=-1
            )
        )
        combined_repr = scatter(
            combined_repr,
            count_to_group_index(neighbor_count),
            dim=0,
            dim_size=len(neighbor_count),
        )

        node_repr = self.mlp_combine_node_nei(
            torch.cat([combined_repr, repr[nodes]], dim=-1)
        )
        return node_repr


class ReprIndexTransform(Pooler):
    def forward(self, repr, ind):
        return repr[ind]


class EmbTransform(Transform):
    def __init__(self, emb_dim, num_embs) -> None:
        super().__init__(emb_dim)
        self.emb = nn.Embedding(num_embs, emb_dim, sparse=False)

    def forward(self, ind):
        return self.emb(ind)


class ScatterReprTransform(Pooler):
    def __init__(self, scatter_method="sum"):
        super().__init__()
        self.scatter_method = scatter_method

    def forward(self, repr, ind, ind_block):
        gd_repr = repr[ind]
        gd_repr = scatter(
            gd_repr,
            count_to_group_index(ind_block),
            dim=0,
            dim_size=len(ind_block),
            reduce=self.scatter_method,
        )
        return gd_repr


class VerGDTransform(Transform):
    """Vertical GD representation for links"""

    def __init__(self, emb_dim, gd_deg=False) -> None:
        super().__init__(emb_dim)
        self.gd_deg = gd_deg
        if gd_deg:
            self.mlp_combine_gd_deg = MLP([emb_dim + 1, 2 * emb_dim, emb_dim])
        self.mlp_gd_process = MLP([emb_dim, 2 * emb_dim, emb_dim])

    def get_ver_gd_one_side(self, repr, gd, gd_len, gd_deg):
        """Get vertical geodesics of one side.

        Arguments:
            repr {torch.tensor} -- graph node representations
            gd {torch.tensor} -- vertical geodesics
            gd_len {torch.tensor} -- count of each vertical geodesics
            gd_deg {torch.tensor} -- degrees of each vertical gd node
        """
        gd_repr = repr[gd]
        if gd_deg:
            gd_repr = self.mlp_combine_gd_deg(
                torch.cat([gd_repr, gd_deg.view(-1, 1)], dim=-1)
            )
        gd_repr = scatter(
            gd_repr,
            count_to_group_index(gd_len),
            dim=0,
            dim_size=len(gd_len),
        )
        return gd_repr

    def forward(
        self,
        repr,
        gd,
        gd_len,
        gd_deg=None,
    ):
        gd_repr = self.get_ver_gd_one_side(repr, gd, gd_len, gd_deg)
        return self.mlp_gd_process(gd_repr)


class ReshapeTransform(Transform):
    def __init__(self, emb_dim):
        super().__init__(emb_dim)

    def forward(self, value):
        return value.view(-1, self.emb_dim)

gp/nn/pooling.py
LechengKong/OneForAll
"""Class Resovler for Efficient Typing
"""

import inspect
from typing import Any, Dict, List, Optional, Union

from torch import Tensor


def normalize_string(s: str) -> str:
    return s.lower().replace("-", "").replace("_", "").replace(" ", "")


def resolver(
    classes: List[Any],
    class_dict: Dict[str, Any],
    query: Union[Any, str],
    base_cls: Optional[Any],
    base_cls_repr: Optional[str],
    *args,
    **kwargs,
):

    if not isinstance(query, str):
        return query

    query_repr = normalize_string(query)
    if base_cls_repr is None:
        base_cls_repr = base_cls.__name__ if base_cls else ""
    base_cls_repr = normalize_string(base_cls_repr)

    for key_repr, cls in class_dict.items():
        if query_repr == key_repr:
            if inspect.isclass(cls):
                obj = cls(*args, **kwargs)
                assert callable(obj)
                return obj
            assert callable(cls)
            return cls

    for cls in classes:
        cls_repr = normalize_string(cls.__name__)
        if query_repr in [cls_repr, cls_repr.replace(base_cls_repr, "")]:
            if inspect.isclass(cls):
                obj = cls(*args, **kwargs)
                assert callable(obj)
                return obj
            assert callable(cls)
            return cls

    choices = set(cls.__name__ for cls in classes) | set(class_dict.keys())
    raise ValueError(f"Could not resolve '{query}' among choices {choices}")


# Activation Resolver #########################################################


def swish(x: Tensor) -> Tensor:
    return x * x.sigmoid()


def activation_resolver(query: Union[Any, str] = "relu", *args, **kwargs):
    import torch

    base_cls = torch.nn.Module
    base_cls_repr = "Act"
    acts = [
        act
        for act in vars(torch.nn.modules.activation).values()
        if isinstance(act, type) and issubclass(act, base_cls)
    ]
    acts += [
        swish,
    ]
    act_dict = {}
    return resolver(
        acts, act_dict, query, base_cls, base_cls_repr, *args, **kwargs
    )

gp/nn/resolver.py
LechengKong/OneForAll
import torch_geometric as pyg
from abc import ABCMeta, abstractmethod

from torch.utils.data import Dataset



class DatasetWithCollate(Dataset, metaclass=ABCMeta):
    @abstractmethod
    def __init__(self) -> None:
        super().__init__()

    @abstractmethod
    def __len__(self):
        pass

    @abstractmethod
    def __getitem__(self, index):
        pass

    @abstractmethod
    def get_collate_fn(self):
        pass


class DGLSingleGraphDataset(DatasetWithCollate):
    def __init__(self, graph):
        super().__init__()
        self.num_nodes = graph.num_nodes()
        self.graph = graph
        self.adj_mat = self.graph.adjacency_matrix(
            transpose=False, scipy_fmt="csr"
        )


class PyGSingleGraphDataset(DatasetWithCollate):
    def __init__(self, graph: pyg.data):
        super().__init__()
        self.num_nodes = graph.num_nodes()
        self.graph = graph
        self.adj_mat = pyg.utils.to_scipy_sparse_matrix(
            graph.edge_index
        ).tocsr()
gp/utils/datasets.py
LechengKong/OneForAll
import numpy as np
import torch
# import dgl

def construct_dgl_graph_from_edges(
    ori_head,
    ori_tail,
    n_entities,
    inverse_edge=False,
    edge_type=None,
    num_rels=None,
):
    num_rels = 1
    if inverse_edge:
        head = np.concatenate([ori_head, ori_tail])
        tail = np.concatenate([ori_tail, ori_head])
    else:
        head = ori_head
        tail = ori_tail
    g = dgl.graph((head, tail), num_nodes=n_entities)
    g.edata["src_node"] = torch.tensor(head, dtype=torch.long)
    g.edata["dst_node"] = torch.tensor(tail, dtype=torch.long)
    if edge_type is not None:
        if num_rels is None:
            num_rels = np.max(edge_type) + 1
        g.edata["type"] = torch.tensor(
            np.concatenate((edge_type, edge_type + num_rels))
        )
    return g


def sample_fixed_hop_size_neighbor(adj_mat, root, hop, max_nodes_per_hope=500):
    visited = np.array(root)
    fringe = np.array(root)
    nodes = np.array([])
    for h in range(1, hop + 1):
        u = adj_mat[fringe].nonzero()[1]
        fringe = np.setdiff1d(u, visited)
        visited = np.union1d(visited, fringe)
        if len(fringe) > max_nodes_per_hope:
            fringe = np.random.choice(fringe, max_nodes_per_hope)
        if len(fringe) == 0:
            break
        nodes = np.concatenate([nodes, fringe])
        # dist_list+=[dist+1]*len(fringe)
    return nodes


def get_k_hop_neighbors(adj_mat, root, hop, block_node=None):
    """Return k-hop neighbor dictionary of root.
    hop2neighbor[i] = the nodes that are exactly i distance away from root.
    """
    if block_node:
        visited = np.array([root, block_node])
    else:
        visited = np.array([root])
    fringe = np.array([root])
    hop2neighbor = {}
    hop2neighbor[0] = fringe
    for h in range(1, hop + 1):
        u = adj_mat[fringe].nonzero()[1]
        fringe = np.setdiff1d(u, visited)
        visited = np.union1d(visited, fringe)
        if len(fringe) == 0:
            break
        hop2neighbor[h] = fringe
        if block_node and h == 1:
            visited = np.setdiff1d(visited, np.array([block_node]))

    return hop2neighbor


def shortest_dist_sparse_mult(adj_mat, hop=6, source=None):
    if source is not None:
        neighbor_adj = adj_mat[source]
        ind = source
    else:
        neighbor_adj = adj_mat
        ind = np.arange(adj_mat.shape[0])
    neighbor_adj_set = [neighbor_adj]
    neighbor_dist = neighbor_adj.todense()
    for i in range(hop - 1):
        new_adj = neighbor_adj_set[i].dot(adj_mat)
        neighbor_adj_set.append(new_adj)
        update_ind = (new_adj.sign() - np.sign(neighbor_dist)) == 1
        r, c = update_ind.nonzero()
        neighbor_dist[r, c] = i + 2
    neighbor_dist[neighbor_dist < 1] = 9999
    neighbor_dist[np.arange(len(neighbor_dist)), ind] = 0
    return np.asarray(neighbor_dist)


def remove_gt_graph_edge(gt_graph, s, t):
    edges = gt_graph.edge(s, t, all_edges=True)
    for e in edges:
        gt_graph.remove_edge(e)
    if gt_graph.is_directed():
        edges = gt_graph.edge(t, s, all_edges=True)
        for e in edges:
            gt_graph.remove_edge(e)


def add_gt_graph_edge(gt_graph, s, t):
    gt_graph.add_edge(s, t)
    if gt_graph.is_directed():
        gt_graph.add_edge(t, s)
gp/utils/graph.py
LechengKong/OneForAll
import numpy as np
from scipy.sparse import csr_matrix
import os.path as osp
import os
import torch
import pickle as pkl
import yaml
from tqdm import tqdm

from gp.utils.graph import construct_dgl_graph_from_edges


def read_knowledge_graph(files, relation2id=None):
    entity2id = {}
    if relation2id is None:
        relation2id = {}

    converted_triplets = {}
    rel_list = [[] for i in range(len(relation2id))]

    ent = 0
    rel = len(relation2id)

    for file_type, file_path in files.items():

        data = []
        with open(file_path) as f:
            file_data = [line.split() for line in f.read().split("\n")[:-1]]

        for triplet in file_data:
            if triplet[0] not in entity2id:
                entity2id[triplet[0]] = ent
                ent += 1
            if triplet[2] not in entity2id:
                entity2id[triplet[2]] = ent
                ent += 1
            if triplet[1] not in relation2id:
                relation2id[triplet[1]] = rel
                rel += 1
                rel_list.append([])

            data.append(
                [
                    entity2id[triplet[0]],
                    relation2id[triplet[1]],
                    entity2id[triplet[2]],
                ]
            )

        for trip in data:
            rel_list[trip[1]].append([trip[0], trip[2]])

        converted_triplets[file_type] = np.array(data)

    id2entity = {v: k for k, v in entity2id.items()}
    id2relation = {v: k for k, v in relation2id.items()}

    adj_list = []
    for rel_mat in rel_list:
        rel_array = np.array(rel_mat)
        if len(rel_array) == 0:
            adj_list.append(csr_matrix((len(entity2id), len(entity2id))))
        else:
            adj_list.append(
                csr_matrix(
                    (
                        np.ones(len(rel_mat)),
                        (rel_array[:, 0], rel_array[:, 1]),
                    ),
                    shape=(len(entity2id), len(entity2id)),
                )
            )

    return (
        adj_list,
        converted_triplets,
        entity2id,
        relation2id,
        id2entity,
        id2relation,
    )


def save_load_torch_data(
    folder_path,
    data,
    num_output=1,
    data_fold=5,
    data_name="saved_gd_data",
    num_workers=32,
):
    saved_data_path = osp.join(folder_path, data_name)
    if not osp.exists(saved_data_path):
        os.mkdir(saved_data_path)
        dt = torch.utils.data.DataLoader(
            data,
            batch_size=256,
            num_workers=32,
        )
        pbar = tqdm(dt)
        fold_len = int(len(dt) / (data_fold - 1))
        count = 0
        fold_count = 0
        for i, t in enumerate(pbar):
            if count == 0:
                data_col = []
                for j in range(num_output):
                    data_col.append([])
            if num_output == 1:
                data_col[0].append(t)
            else:
                for j, v in enumerate(t):
                    data_col[j].append(v)
            count += 1
            if count == fold_len:
                for i, it in enumerate(data_col):
                    cdata = torch.cat(it, dim=0).numpy()
                    np.save(
                        osp.join(
                            saved_data_path, str(i) + "_" + str(fold_count)
                        ),
                        cdata,
                    )
                    for itm in it:
                        del itm
                    del cdata
                fold_count += 1
                count = 0
        if count > 0:
            for i, it in enumerate(data_col):
                cdata = torch.cat(it, dim=0).numpy()
                np.save(
                    osp.join(saved_data_path, str(i) + "_" + str(fold_count)),
                    cdata,
                )
    saved_data = []
    if num_output == 1:
        for j in range(data_fold):
            ipath = osp.join(saved_data_path, str(0) + "_" + str(j) + ".npy")
            if osp.exists(ipath):
                saved_data.append(np.load(ipath))
    else:
        for i in range(num_output):
            saved_data.append([])
            for j in range(data_fold):
                ipath = osp.join(
                    saved_data_path, str(i) + "_" + str(j) + ".npy"
                )
                if osp.exists(ipath):
                    saved_data[i].append(np.load(ipath))
    return saved_data, data_fold


def load_exp_dataset_dgl(directory):
    graphs = []
    labels = []
    with open(directory, "r") as data:
        num_graphs = int(data.readline().rstrip().split(" ")[0])
        for i in range(num_graphs):
            graph_meta = data.readline().rstrip().split(" ")
            num_vertex = int(graph_meta[0])
            # curr_graph = np.zeros(shape=(num_vertex, num_vertex))
            labels.append(int(graph_meta[1]))
            node_labels = np.zeros((num_vertex, 2))
            edges = []
            for j in range(num_vertex):
                vertex = data.readline().rstrip().split(" ")
                node_labels[j, int(vertex[0])] = 1
                for k in range(2, len(vertex)):
                    edges.append([j, int(vertex[k])])
            edges = np.array(edges)
            g = construct_dgl_graph_from_edges(
                edges[:, 0],
                edges[:, 1],
                n_entities=num_vertex,
                inverse_edge=True,
            )
            g.ndata["feat"] = torch.tensor(node_labels, dtype=torch.float)
            graphs.append(g)
    return graphs, np.array(labels)


def open_and_load_pickle(filename):
    open_file = open(filename, "rb")
    data = pkl.load(open_file)
    open_file.close()
    return data


def load_yaml(dir):
    with open(dir, "r") as stream:
        return yaml.safe_load(stream)

gp/utils/io.py
LechengKong/OneForAll
from itertools import product
import time
import numpy as np
from sklearn.model_selection import StratifiedKFold
import torch
import random
import os
import json
import yaml

import os.path as osp
from datetime import datetime
from gp.utils.io import load_yaml
from types import SimpleNamespace


class SmartTimer:
    """A timer utility that output the elapsed time between this
    call and last call.
    """

    def __init__(self, verb=True) -> None:
        """SmartTimer Constructor

        Keyword Arguments:
            verb {bool} -- Controls printing of the timer (default: {True})
        """
        self.last = time.time()
        self.verb = verb

    def record(self):
        """Record current timestamp"""
        self.last = time.time()

    def cal_and_update(self, name):
        """Record current timestamp and print out time elapsed from last
        recorded time.

        Arguments:
            name {string} -- identifier of the printout.
        """
        now = time.time()
        if self.verb:
            print(name, now - self.last)
        self.record()


class SparseData:
    def __init__(self, data, data_count=None, data_offset=None):
        self.data = data
        if data_count is None and data_offset is None:
            if isinstance(data, list):
                self.num_data = len(data)
                self.data_count = np.array([len(d) for d in data])
                self.data = np.concatenate(data, axis=0)
            elif isinstance(data, np.ndarray):
                self.num_data = 1
                self.data_count = len(data)
            self.data_offset = self.count2offset(self.data_count)
        if self.data_count is None:
            self.data_count = self.offset2count(self.data_offset)
        if self.data_offset is None:
            self.data_offset = self.count2offset(self.data_count)

    def count2offset(self, count):
        return np.r_[0, np.cumsum(count[:-1])]

    def offset2count(self, offset):
        return np.r_[offset[1:], len(self.data)] - offset


def save_params(filename, params):
    """Write a Namespace object to file

    Arguments:
        filename {string} -- destination of the saved file
        params {Namespace} -- namespace object
    """
    d = vars(params)
    with open(filename, "a") as f:
        json.dump(d, f)


def set_random_seed(seed):
    """Set python, numpy, pytorch global random seed.
    Does not guarantee determinism due to PyTorch's feature.

    Arguments:
        seed {int} -- Random seed to set
    """
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    np.random.seed(seed)
    random.seed(seed)
    os.environ["PYTHONHASHSEED"] = str(seed)


def sparse_uniform_select(data, data_count, p=0.5):
    """Randomly select from a sparse representation.

    Arguments:
        data {numpy.ndarray} -- Sparse data
        data_count {numpy.ndarray} -- Sparse index of the data

    Keyword Arguments:
        p {int or numpy.ndarray} -- If int, mu=p for all entries,
        if numpy.ndarray, should be of the same shape as data_count.
         (default: {0.5})

    Returns:
        selected_data -- selected_data
    """
    if isinstance(p, np.ndarray):
        p = np.repeat(p, data_count)
    prob = np.random.rand(len(data))
    select = prob < p
    data_index = np.arange(len(data_count)).repeat(data_count)
    new_data_count = np.bincount(data_index[select], minlength=len(data_count))
    return data[select], new_data_count


def sparse_uniform_sample(data, data_count, c=1):
    """
    Return a DGL graph specified by edges.

    :param head: Edge head.
    :param tail: Edge tail.
    :return: DGLGraph.
    :rtype: DGLGraph
    """
    if isinstance(c, int):
        c = np.repeat(c, len(data_count))
        c = (data_count > 0) * c
    max_val = np.max(data_count) * 10
    select_ind = np.random.randint(max_val, size=np.sum(c))
    select_ind = select_ind % np.repeat(data_count, c)
    offset = np.r_[0, data_count[:-1]]
    offset = np.cumsum(offset)
    select_ind = select_ind + offset.repeat(c)
    return data[select_ind], c


def k_fold_ind(labels, fold):
    """Generate stratified k fold split index based on labels

    Arguments:
        labels {np.ndarray} -- labels of the data
        fold {int} -- number of folds

    Returns:
        list[numpy.ndarray] -- A list whose elements are indices of data
        in the fold.
    """
    ksfold = StratifiedKFold(n_splits=fold, shuffle=True, random_state=10)
    folds = []
    for _, t_index in ksfold.split(
        np.zeros_like(np.array(labels)), np.array(labels, dtype=int)
    ):
        folds.append(t_index)
    return folds


def k_fold2_split(folds, data_len):
    """Split the data index into train/test/validation based on fold,
    one fold for testing, one fold for validation and the rest for training.

    Arguments:
        folds {list[numpy.ndarray]} -- fold information
        data_len {int} -- lenght of the data

    Returns:
        list[list[numpy.ndarray]] -- a list of train/test/validation split
        indices.
    """
    splits = []
    for i in range(len(folds)):
        test_arr = np.zeros(data_len, dtype=bool)
        test_arr[folds[i]] = 1
        val_arr = np.zeros(data_len, dtype=bool)
        val_arr[folds[int((i + 1) % len(folds))]] = 1
        train_arr = np.logical_not(np.logical_or(test_arr, val_arr))
        train_ind = train_arr.nonzero()[0]
        test_ind = test_arr.nonzero()[0]
        val_ind = val_arr.nonzero()[0]
        splits.append([train_ind, test_ind, val_ind])
    return splits


def dict_res_summary(res_col):
    """Combine multiple dictionary information into one dictionary
    so that all entries with the same key will be concatenated into
    a list

    Arguments:
        res_col {list[dictionary]} -- a list of dictionary

    Returns:
        dictionary -- summarized dictionary information
    """
    res_dict = {}
    for res in res_col:
        for k in res:
            if k not in res_dict:
                res_dict[k] = []
            res_dict[k].append(res[k])
    return res_dict


def multi_data_average_exp(data, args, exp):
    val_res_col = []
    test_res_col = []
    for split in data:
        val_res, test_res = exp(split, args)
        val_res_col.append(val_res)
        test_res_col.append(test_res)

    val_res_dict = dict_res_summary(val_res_col)
    test_res_dict = dict_res_summary(test_res_col)
    return val_res_dict, test_res_dict


def hyperparameter_grid_search(
    hparams, data, exp, args, search_metric, evaluator, exp_arg=None
):
    named_params = [[(k, p) for p in hparams[k]] for k in hparams]
    best_met = evaluator.init_result()
    best_res = None
    params = product(*named_params)
    for p in params:
        for name, val in p:
            setattr(args, name, val)
        if exp_arg:
            val_res, test_res = exp(data, args, exp_arg)
        else:
            val_res, test_res = exp(data, args)
        val_metric_res, test_metric_res = np.array(
            val_res[search_metric]
        ), np.array(test_res[search_metric])
        val_mean, val_std = np.mean(val_metric_res), np.std(val_metric_res)
        test_mean, test_std = np.mean(test_metric_res), np.std(test_metric_res)
        if evaluator.better_results(val_mean, best_met):
            best_met = val_mean
            best_res = {
                "metric": search_metric,
                "val_mean": val_mean,
                "val_std": val_std,
                "test_mean": test_mean,
                "test_std": test_std,
                "full_val": val_res,
                "full_test": test_res,
                "params": p,
            }
    return best_res


def write_res_to_file(
    file,
    dataset,
    metric,
    test_mean,
    val_mean=0,
    test_std=0,
    val_std=0,
    params=None,
    res=None,
):
    with open(file, "a") as f:
        f.write("\n\n")
        f.write(res)
        f.write("\n")
        f.write("Dataset: {} \n".format(dataset))
        f.write("Optimize wrt {}\n".format(metric))
        f.write("val, {:.5f} ± {:.5f} \n".format(val_mean, val_std))
        f.write("test, {:.5f} ± {:.5f} \n".format(test_mean, test_std))
        f.write("best res:")
        f.write(str(params))


def var_size_repeat(size, chunks, repeats):
    a = np.arange(size)
    s = np.r_[0, chunks.cumsum()]
    starts = a[np.repeat(s[:-1], repeats)]
    if len(starts) == 0:
        return np.array([], dtype=int)
    chunk_rep = np.repeat(chunks, repeats)
    ends = starts + chunk_rep

    clens = chunk_rep.cumsum()
    ids = np.ones(clens[-1], dtype=int)
    ids[0] = starts[0]
    ids[clens[:-1]] = starts[1:] - ends[:-1] + 1
    out = ids.cumsum()
    return out


def count_to_group_index(count):
    return torch.arange(len(count), device=count.device).repeat_interleave(
        count
    )


def setup_exp(params):
    if not osp.exists("./saved_exp"):
        os.mkdir("./saved_exp")

    curtime = datetime.now()
    exp_dir = osp.join("./saved_exp", str(curtime))
    os.mkdir(exp_dir)
    with open(osp.join(exp_dir, "command"), "w") as f:
        yaml.dump(params, f)
    params["exp_dir"] = exp_dir


def combine_dict(*args):
    combined_dict = {}
    for d in args:
        for k in d:
            combined_dict[k] = d[k]
    return combined_dict


def merge_mod(params, mod_args):
    for i in range(0, len(mod_args), 2):
        if mod_args[i + 1].isdigit():
            val = int(mod_args[i + 1])
        elif mod_args[i + 1].replace(".", "", 1).isdigit():
            val = float(mod_args[i + 1])
        elif mod_args[i + 1].lower() == "true":
            val = True
        elif mod_args[i + 1].lower() == "false":
            val = False
        else:
            val = mod_args[i + 1]
        params[mod_args[i]] = val
    return params


def convert_yaml_params(params_path):
    load_params = load_yaml(params_path)
    load_params = SimpleNamespace(**load_params)
    return load_params

gp/utils/utils.py
LechengKong/OneForAll
import torch
from gp.lightning.module_template import BaseTemplate
from lightning.pytorch import LightningModule
import os.path as osp


class GraphPredLightning(BaseTemplate):
    def forward(self, batch):
        return self.model(batch)


class TokenPredLightning(BaseTemplate):
    def forward(self, batch):
        return self.model(batch)

    def validation_step(self, batch, batch_idx):
        tokens = self.model.generate(batch)
        with torch.no_grad():
            if self.eval_kit.has_eval_state("valid"):
                self.eval_kit.eval_step(tokens, batch, "valid")
        self.val_names = ["valid"]

    def test_step(self, batch, batch_idx):
        tokens = self.model.generate(batch)
        with torch.no_grad():
            if self.eval_kit.has_eval_state("test"):
                self.eval_kit.eval_step(tokens, batch, "test")


class GraphFinetuneLightning(BaseTemplate):
    def forward(self, batch):
        return self.model(batch)

    def configure_optimizers(self):
        optimizer = self.exp_config.optimizer(
            self.exp_config.opt_params
            if self.exp_config.opt_params is not None
            else filter(lambda p: p.requires_grad, self.parameters()),
            **self.exp_config.optimizer_args
        )
        return {
            "optimizer": optimizer,
        }


class CLBaseTemplate(LightningModule):
    def __init__(
        self,
        exp_config,
        encode_model: torch.nn.Module,
        loss,
        params,
        eval_params,
        proj_head=None,
        pred_head=None,
        sup_cl=False,
        name = "",
    ):

        super().__init__()

        self.exp_config = exp_config
        self.model = encode_model
        self.loss = loss
        self.params = params
        self.eval_params = eval_params
        self.name = name
        self.proj_head = proj_head
        self.pred_head = pred_head
        self.sup_cl = sup_cl
        if self.pred_head is not None:
            self.bin_loss = torch.nn.BCEWithLogitsLoss()

    def forward(self, batch):
        return self.model(batch)

    def configure_optimizers(self):
        optimizer = self.exp_config.optimizer(
            self.exp_config.opt_params
            if self.exp_config.opt_params is not None
            else self.model.parameters(),
            **self.exp_config.optimizer_args
        )
        return {
            "optimizer": optimizer,
        }

    def _calculate_cl_loss(self, g1, g2, z1, z2):

        # mask1 = g1.h_node_mask + g1.true_nodes_mask + g1.spt_nodes_mask + g1.target_node_mask
        # mask2 = g2.h_node_mask + g2.true_nodes_mask + g2.spt_nodes_mask + g2.target_node_mask
        mask1 = g1.true_nodes_mask
        mask2 = g2.true_nodes_mask
        if self.proj_head != None:
            z1 = self.proj_head(z1)[mask1]
            z2 = self.proj_head(z2)[mask2]

        z1 = torch.nn.functional.normalize(z1, dim=1)
        z2 = torch.nn.functional.normalize(z2, dim=1)
        h_node_embs = torch.stack([z1, z2], dim=1)

        # generate supervised mask for h_nodes
        # mask = torch.arange(self.params.n_way).repeat_interleave(self.params.q_query)

        # from losses import SupConLoss
        # loss = SupConLoss(temperature=0.07)
        #sup_loss = loss(features=h_node_embs, labels=mask)
        if self.sup_cl:
            print(torch.matmul(g1.bin_labels.view(-1,1), g1.bin_labels.view(1,-1)))
            cl_loss = self.loss(features=h_node_embs, mask=torch.matmul(g1.bin_labels.view(-1,1), g1.bin_labels.view(1,-1)))
        else:
            cl_loss = self.loss(features=h_node_embs, )
        return cl_loss

    def _calculate_binary_cls_loss(self, g, z):
        z = self.pred_head(z)
        score = z[g.true_nodes_mask]
        labels = g.bin_labels
        valid_ind = labels == labels
        return self.bin_loss(score.view(-1)[valid_ind], labels[valid_ind])

    def combine_loss(self, batch):
        g1, g2 = batch
        z1 = self(g1)
        z2 = self(g2)
        cl_loss = self._calculate_cl_loss(g1, g2, z1, z2)
        if self.pred_head is not None:
            binary_cls_loss = self._calculate_binary_cls_loss(g1, z1) + self._calculate_binary_cls_loss(g2, z2)
            return (cl_loss, cl_loss+binary_cls_loss, binary_cls_loss,)
        return (cl_loss, cl_loss)


    def training_step(self, batch, batch_idx):
        # get embeddings of h nodes (n_way x q_query x fs_task_num, emb_dim)
        # current only consider fs_task_num = 1
        loss = self.combine_loss(batch)
        loss_name = [osp.join(self.name, "train", "cl_loss"),
                     osp.join(self.name, "train", "loss"),
                     osp.join(self.name, "train", "pred_loss"),]
        log_dict = {loss_name[idx]: loss[idx] for idx in range(len(loss))}

        self.log_dict(
            log_dict,
            on_step=True,
            on_epoch=True,
            prog_bar=True,
            batch_size=batch.batch_size
            if hasattr(batch, "batch_size")
            else len(batch),
        )
        return loss[1]

    def validation_step(self, batch, batch_idx):
        # get embeddings of h nodes (n_way x q_query x fs_task_num, emb_dim)
        # current only consider fs_task_num = 1
        loss = self.combine_loss(batch)
        loss_name = [osp.join(self.name, "val", "cl_loss"),
                     osp.join(self.name, "val", "loss"),
                     osp.join(self.name, "val", "pred_loss")]
        log_dict = {loss_name[idx]: loss[idx] for idx in range(len(loss))}

        self.log_dict(
            log_dict,
            on_step=True,
            on_epoch=True,
            prog_bar=True,
            batch_size=batch.batch_size
            if hasattr(batch, "batch_size")
            else len(batch),
        )
        return loss[1]


class SimPredLightning(LightningModule):
    def __init__(
        self,
        exp_config,
        encode_model: torch.nn.Module,
        loss,
        head = None,
    ):

        super().__init__()

        self.exp_config = exp_config
        self.model = encode_model
        self.loss = loss
        self.head = head

    def forward(self, batch):
        if self.head != None:
            return self.head(self.model(batch))
        else:
            return self.model(batch)

    def configure_optimizers(self):
        optimizer = self.exp_config.optimizer(
            self.exp_config.opt_params
            if self.exp_config.opt_params is not None
            else self.model.parameters(),
            **self.exp_config.optimizer_args
        )
        return {
            "optimizer": optimizer,
        }

    def training_step(self, batch, batch_idx):
        # emb = self.batch
        # qry_emb = self.batch[batch.target_node_mask]
        print(batch.x[batch.target_node_mask])
        print(batch.x[batch.true_nodes_mask])
        print('-'*40)
        loss, acc = self.loss(self(batch)[batch.target_node_mask], self(batch)[batch.true_nodes_mask], batch.bin_labels)

        log = {"train_loss": loss, "train_acc": acc}

        self.log_dict(
            log,
            on_step=True,
            on_epoch=True,
            prog_bar=True,
            batch_size=batch.batch_size
            if hasattr(batch, "batch_size")
            else len(batch),
        )

        return loss

    def validation_step(self, batch, batch_idx, dataloader_idx):
        loss, acc = self.loss(self(batch)[batch.target_node_mask], self(batch)[batch.true_nodes_mask], batch.bin_labels)

        if dataloader_idx == 0:
            log = {"val_loss": loss, "val_acc": acc}
        elif dataloader_idx == 1:
            log = {"t_val_loss": loss, "t_val_acc": acc}

        log = {"idk_loss": loss, "idk_acc": acc}

        self.log_dict(
            log,
            on_step=False,
            on_epoch=True,
            prog_bar=True,
            batch_size=batch.batch_size
            if hasattr(batch, "batch_size")
            else len(batch),
        )

        return loss


    def test_step(self, batch, batch_idx):
        loss, acc = self.loss(self(batch)[batch.target_node_mask], self(batch)[batch.true_nodes_mask], batch.bin_labels)

        log = {"test_loss": loss, "test_acc": acc}

        self.log_dict(
            log,
            on_step=False,
            on_epoch=True,
            prog_bar=True,
            batch_size=batch.batch_size
            if hasattr(batch, "batch_size")
            else len(batch),
        )

        return loss


lightning_model.py
LechengKong/OneForAll
import torch
from transformers import LlamaTokenizer, LlamaForCausalLM

## v2 models
model_path = "openlm-research/open_llama_3b_v2"

## v1 models
# model_path = 'openlm-research/open_llama_3b'
# model_path = 'openlm-research/open_llama_7b'
# model_path = 'openlm-research/open_llama_13b'

tokenizer = LlamaTokenizer.from_pretrained(model_path)
model = LlamaForCausalLM.from_pretrained(
    model_path,
    torch_dtype=torch.float16,
    device_map="auto",
)

prompt = "Q: Among machine learning, biology, and chemistry, which one is computer science most pertained to?\nA:"
input_ids = tokenizer(prompt, return_tensors="pt").input_ids

generation_output = model.generate(
    input_ids=input_ids.to("cuda:0"), max_new_tokens=32
)
print(tokenizer.decode(generation_output[0]))
print(torch.cuda.max_memory_allocated("cuda:0"))

llama_exp.py
LechengKong/OneForAll
from __future__ import print_function

import torch
import torch.nn as nn
from torch.nn import functional as F
from torch.nn.modules import Module


class SupConLoss(Module):
    """Supervised Contrastive Learning: https://arxiv.org/pdf/2004.11362.pdf.
    It also supports the unsupervised contrastive loss in SimCLR"""
    def __init__(self, temperature=0.07, contrast_mode='all',
                 base_temperature=0.07):
        super(SupConLoss, self).__init__()
        self.temperature = temperature
        self.contrast_mode = contrast_mode
        self.base_temperature = base_temperature

    def forward(self, features, labels=None, mask=None):
        """Compute loss for model. If both `labels` and `mask` are None,
        it degenerates to SimCLR unsupervised loss:
        https://arxiv.org/pdf/2002.05709.pdf

        Args:
            features: hidden vector of shape [bsz, n_views, ...].
            labels: ground truth of shape [bsz].
            mask: contrastive mask of shape [bsz, bsz], mask_{i,j}=1 if sample j
                has the same class as sample i. Can be asymmetric.
        Returns:
            A loss scalar.
        """
        device = (torch.device('cuda')
                  if features.is_cuda
                  else torch.device('cpu'))

        if len(features.shape) < 3:
            raise ValueError('`features` needs to be [bsz, n_views, ...],'
                             'at least 3 dimensions are required')
        if len(features.shape) > 3:
            features = features.view(features.shape[0], features.shape[1], -1)

        batch_size = features.shape[0]
        if labels is not None and mask is not None:
            raise ValueError('Cannot define both `labels` and `mask`')
        elif labels is None and mask is None:
            mask = torch.eye(batch_size, dtype=torch.float32).to(device)
        elif labels is not None:
            labels = labels.contiguous().view(-1, 1)
            if labels.shape[0] != batch_size:
                raise ValueError('Num of labels does not match num of features')
            mask = torch.eq(labels, labels.T).float().to(device)
        else:
            mask = mask.float().to(device)

        contrast_count = features.shape[1]
        contrast_feature = torch.cat(torch.unbind(features, dim=1), dim=0)
        if self.contrast_mode == 'one':
            anchor_feature = features[:, 0]
            anchor_count = 1
        elif self.contrast_mode == 'all':
            anchor_feature = contrast_feature
            anchor_count = contrast_count
        else:
            raise ValueError('Unknown mode: {}'.format(self.contrast_mode))

        # compute logits
        anchor_dot_contrast = torch.div(
            torch.matmul(anchor_feature, contrast_feature.T),
            self.temperature)
        # for numerical stability
        logits_max, _ = torch.max(anchor_dot_contrast, dim=1, keepdim=True)
        logits = anchor_dot_contrast - logits_max.detach()

        # tile mask
        mask = mask.repeat(anchor_count, contrast_count)
        # mask-out self-contrast cases
        logits_mask = torch.scatter(
            torch.ones_like(mask),
            1,
            torch.arange(batch_size * anchor_count).view(-1, 1).to(device),
            0
        )
        mask = mask * logits_mask

        # compute log_prob
        exp_logits = torch.exp(logits) * logits_mask
        log_prob = logits - torch.log(exp_logits.sum(1, keepdim=True))

        # compute mean of log-likelihood over positive
        mean_log_prob_pos = (mask * log_prob).sum(1) / mask.sum(1)

        # loss
        loss = - (self.temperature / self.base_temperature) * mean_log_prob_pos
        loss = loss.view(anchor_count, batch_size).mean()

        return loss



class PrototypicalLoss(Module):
    '''
    Loss class deriving from Module for the prototypical loss function defined below
    '''
    def __init__(self, n_way):
        super(PrototypicalLoss, self).__init__()
        self.n_way = n_way

    def euclidean_dist(self, x, y):
        '''
            Compute euclidean distance between two tensors
            '''
        # x: N x D
        # y: M x D
        n = x.size(0)
        m = y.size(0)
        d = x.size(1)
        if d != y.size(1):
            raise Exception
        x = x.unsqueeze(1).expand(n, m, d)
        y = y.unsqueeze(0).expand(n, m, d)
        return torch.pow(x - y, 2).sum(2)

    def forward(self, qry_emb, prompt_emb, bin_labels):
        print(qry_emb)
        print(prompt_emb)
        dists = self.euclidean_dist(qry_emb, prompt_emb)
        # Extract required slices based on your loop logic
        sliced_dists = [dists[i, i * self.n_way: (i + 1) * self.n_way] for i in range(qry_emb.size(0))]


        concat_dists = torch.stack(sliced_dists, dim=0)
        # Compute loss using a vectorized approach
        log_probs = F.log_softmax(-concat_dists, dim=1)
        loss = -(log_probs * bin_labels.view(-1,self.n_way)).sum() / qry_emb.size(0)

        _, y_hat = log_probs.max(1)
        # print(log_probs)
        # print(y_hat)
        #print(bin_labels.view(-1,self.n_way))
        acc_val = y_hat.eq(torch.argmax(bin_labels, dim=-1)).float().mean()

        return loss, acc_val



# def euclidean_dist(x, y):
#     n = x.size(0)
#     m = y.size(0)
#     d = x.size(1)
#     if d != y.size(1):
#         raise Exception
#     x = x.unsqueeze(1).expand(n, m, d)
#     y = y.unsqueeze(0).expand(n, m, d)
#     return torch.pow(x - y, 2).sum(2)

# for i in range(a.size(0)):
#     aa = a[i,:].view(1,-1)
#     bb = b[i * n_way:(i+1) *n_way, :]
#     dd = euclidean_dist(aa, bb)
#     loss = (F.log_softmax(-dd, dim=1) * bin_labels[i * n_way: (i+1) * n_way]).sum()
#     losses += loss
#     _, y_hat = log_p_y.max(2)
#     acc_val = y_hat.eq(target_inds.squeeze(2)).float().mean()
#
#     return loss_val, acc_val
#
#
# def prototypical_loss(qry_emb, bin_labels, prompt_emb):
#     '''
#     Inspired by https://github.com/jakesnell/prototypical-networks/blob/master/protonets/models/few_shot.py
#
#     Compute the barycentres by averaging the features of n_support
#     samples for each class in target, computes then the distances from each
#     samples' features to each one of the barycentres, computes the
#     log_probability for each n_query samples for each one of the current
#     classes, of appartaining to a class c, loss and accuracy are then computed
#     and returned
#     Args:
#     - input: the model output for a batch of samples
#     - target: ground truth for the above batch of samples
#     - n_support: number of samples to keep in account when computing
#       barycentres, for each one of the current classes
#     '''
#     target_cpu = target.to('cpu')
#     input_cpu = input.to('cpu')
#
#     def supp_idxs(c):
#         # FIXME when torch will support where as np
#         return target_cpu.eq(c).nonzero()[:n_support].squeeze(1)
#
#     # FIXME when torch.unique will be available on cuda too
#     classes = torch.unique(target_cpu)
#     n_classes = len(classes)
#     # FIXME when torch will support where as np
#     # assuming n_query, n_target constants
#     n_query = target_cpu.eq(classes[0].item()).sum().item() - n_support
#
#     support_idxs = list(map(supp_idxs, classes))
#
#     prototypes = torch.stack([input_cpu[idx_list].mean(0) for idx_list in support_idxs])
#     # FIXME when torch will support where as np
#     query_idxs = torch.stack(list(map(lambda c: target_cpu.eq(c).nonzero()[n_support:], classes))).view(-1)
#
#     query_samples = input.to('cpu')[query_idxs]
#     dists = euclidean_dist(query_samples, prototypes)
#
#     log_p_y = F.log_softmax(-dists, dim=1).view(n_classes, n_query, -1)
#
#     target_inds = torch.arange(0, n_classes)
#     target_inds = target_inds.view(n_classes, 1, 1)
#     target_inds = target_inds.expand(n_classes, n_query, 1).long()
#
#     loss_val = -log_p_y.gather(2, target_inds).squeeze().view(-1).mean()
#     _, y_hat = log_p_y.max(2)
#     acc_val = y_hat.eq(target_inds.squeeze(2)).float().mean()
#
#     return loss_val,  acc_val
losses.py
LechengKong/OneForAll
import torch
from torch import Tensor
from torch import nn

from gp.nn.models.util_model import MLP
from gp.nn.models.GNN import MultiLayerMessagePassing
from gp.nn.layer.pyg import RGCNEdgeConv, RGATEdgeConv
from torch_geometric.nn.pool import global_add_pool

from torch_geometric.transforms.add_positional_encoding import AddRandomWalkPE


class TextClassModel(torch.nn.Module):
    def __init__(self, model, outdim, task_dim, emb=None):
        super().__init__()
        self.model = model
        if emb is not None:
            self.emb = torch.nn.Parameter(emb.clone())

        self.mlp = MLP([2 * outdim, 2 * outdim, outdim, task_dim])

    def forward(self, g):
        emb = self.model(g)
        class_emb = emb[g.target_node_mask]
        att_emb = class_emb.repeat_interleave(len(self.emb), dim=0)
        att_emb = torch.cat(
            [att_emb, self.emb.repeat(len(class_emb), 1)], dim=-1
        )
        res = self.mlp(att_emb).view(-1, len(self.emb))
        return res


class AdaPoolClassModel(torch.nn.Module):
    def __init__(self, model, outdim, task_dim, emb=None):
        super().__init__()
        self.model = model
        if emb is not None:
            self.emb = torch.nn.Parameter(emb.clone())

        self.mlp = MLP([2 * outdim, 2 * outdim, outdim, task_dim])

    def forward(self, g):
        emb = self.model(g)
        float_mask = g.target_node_mask.to(torch.float)
        target_emb = float_mask.view(-1, 1) * emb
        n_count = global_add_pool(float_mask, g.batch, g.num_graphs)
        class_emb = global_add_pool(target_emb, g.batch, g.num_graphs)
        class_emb = class_emb / n_count.view(-1, 1)
        rep_class_emb = class_emb.repeat_interleave(g.num_classes, dim=0)
        res = self.mlp(
            torch.cat([rep_class_emb, g.x[g.true_nodes_mask]], dim=-1)
        )
        return res


class SingleHeadAtt(torch.nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.sqrt_dim = torch.sqrt(torch.tensor(dim))
        self.Wk = torch.nn.Parameter(torch.zeros((dim, dim)))
        torch.nn.init.xavier_uniform_(self.Wk)
        self.Wq = torch.nn.Parameter(torch.zeros((dim, dim)))
        torch.nn.init.xavier_uniform_(self.Wq)

    def forward(self, key, query, value):
        score = torch.bmm(query, key.transpose(1, 2)) / self.sqrt_dim
        attn = torch.nn.functional.softmax(score, -1)
        context = torch.bmm(attn, value)
        return context, attn


class BinGraphModel(torch.nn.Module):
    def __init__(self, model, outdim, task_dim, add_rwpe=None):
        super().__init__()
        self.model = model
        self.mlp = MLP([outdim, 2 * outdim, outdim, task_dim])
        if add_rwpe is not None:
            self.rwpe = AddRandomWalkPE(add_rwpe)
            self.edge_rwpe_prior = torch.nn.Parameter(
                torch.zeros((1, add_rwpe))
            )
            torch.nn.init.xavier_uniform_(self.edge_rwpe_prior)
            self.rwpe_normalization = torch.nn.BatchNorm1d(add_rwpe)
            self.walk_length = add_rwpe
        else:
            self.rwpe = None

    def forward(self, g):
        if self.rwpe is not None:
            with torch.no_grad():
                rwpe_norm = self.rwpe_normalization(g.rwpe)
                g.x = torch.cat([g.x, rwpe_norm], dim=-1)
                g.edge_attr = torch.cat(
                    [
                        g.edge_attr,
                        self.edge_rwpe_prior.repeat(len(g.edge_attr), 1),
                    ],
                    dim=-1,
                )
        emb = self.model(g)
        class_emb = emb[g.true_nodes_mask]
        # print(class_emb)
        res = self.mlp(class_emb)
        return res


class BinGraphAttModel(torch.nn.Module):
    def __init__(self, model, outdim, task_dim, add_rwpe=None):
        super().__init__()
        self.model = model
        self.mlp = MLP([outdim, 2 * outdim, outdim, task_dim])
        self.att = SingleHeadAtt(outdim)
        if add_rwpe is not None:
            self.rwpe = AddRandomWalkPE(add_rwpe)
            self.edge_rwpe_prior = torch.nn.Parameter(
                torch.zeros((1, add_rwpe))
            )
            torch.nn.init.xavier_uniform_(self.edge_rwpe_prior)
            self.rwpe_normalization = torch.nn.BatchNorm1d(add_rwpe)
            self.walk_length = add_rwpe
        else:
            self.rwpe = None

    def forward(self, g):
        if self.rwpe is not None:
            with torch.no_grad():
                rwpe_norm = self.rwpe_normalization(g.rwpe)
                g.x = torch.cat([g.x, rwpe_norm], dim=-1)
                g.edge_attr = torch.cat(
                    [
                        g.edge_attr,
                        self.edge_rwpe_prior.repeat(len(g.edge_attr), 1),
                    ],
                    dim=-1,
                )
        emb = torch.stack(self.model(g), dim=1)
        query = g.x.unsqueeze(1)
        emb = self.att(emb, query, emb)[0].squeeze()

        class_emb = emb[g.true_nodes_mask]
        # print(class_emb)
        res = self.mlp(class_emb)
        return res


class NodeGenerationModel(torch.nn.Module):
    """Model for generating node given prompt node instruction. The input graph with prompt node will be first
    input to a gnn model to learn the topology and semantic information in the graph. Then, the output embedding
    of the prompt node will be input to a transformer model to generate nodes based on the information learned from
    graph and the instruction of the prompt node.
    Args:
        gnn_model (torch.nn.Module): The GNN model.
        out_dim (int): Output dimension of the model.
    """

    def __init__(
        self,
        gnn_model: torch.nn.Module,
        true_emb: Tensor,
        emb_dim: int,
        out_dim: int,
    ):
        super().__init__()
        self.gnn_model = gnn_model
        self.class_mlp = nn.Linear(emb_dim, out_dim)
        self.embedding = nn.Embedding.from_pretrained(true_emb, freeze=True)
        self.transformer = TransformerModel(emb_dim, 6, emb_dim, 8)

    def forward(self, g):
        seq_len = g.y.size(-1)  # B * S
        # message passing
        emb = self.gnn_model(g)
        # get prompt node embedding
        class_emb = emb[g.prompt_idx.bool()]
        batch_size, dim = class_emb.size()

        # construct input for transformer model
        transformer_input = torch.zeros(
            (batch_size, seq_len + 1, dim), device=class_emb.device
        )
        true_emb = self.embedding(g.y.long())
        transformer_input[:, 0, :] = class_emb
        transformer_input[:, 1:, :] = true_emb
        # src_key_padding_mask = torch.cat([torch.ones([batch_size, 1], device=class_emb.device),
        #                                   g.y], dim=-1) == 0
        src_mask = self.get_mask(seq_len + 1).to(class_emb.device)
        t_output = self.transformer(
            transformer_input, mask=src_mask
        )  # B * (S + 1) * H
        next_token_distribution = self.class_mlp(t_output)[:, :-1]
        # print(g.y[:2])
        # print(torch.argmax(next_token_distribution, dim=-1)[:2])
        # print("*" * 89)
        return next_token_distribution.contiguous().view(
            batch_size * seq_len, -1
        )  # B * (S + 1) * D

    def get_mask(self, seq_len):
        mask = torch.triu(
            torch.full((seq_len, seq_len), float("-inf")), diagonal=1
        )
        return mask

    def generate(self, g):
        seq_len = g.y.size(-1)
        # message passing
        emb = self.gnn_model(g)
        # get prompt node embedding
        class_emb = emb[g.prompt_idx.bool()]
        batch_size, dim = class_emb.size()

        # construct input for transformer model
        transformer_input = torch.zeros(
            (batch_size, seq_len + 1, dim), device=class_emb.device
        )
        transformer_input[:, 0, :] = class_emb
        # src_key_padding_mask = torch.cat([torch.ones([batch_size, 1], device=class_emb.device),
        #                                   g.y], dim=-1) == 0
        src_mask = self.get_mask(seq_len + 1).to(class_emb.device)
        all_tokens = []
        for i in range(seq_len):
            # src_key_generation_mask = torch.zeros_like(src_key_padding_mask)
            # src_key_generation_mask[:, : i+1] = 1
            # src_key_mask = torch.logical_or(src_key_padding_mask, src_key_generation_mask == 0)

            t_output = self.transformer(transformer_input, mask=src_mask)

            cur_embedding = t_output[:, i]
            # print(cur_embedding)
            next_token = torch.argmax(
                self.class_mlp(cur_embedding), dim=-1
            )  # B
            all_tokens.append(next_token.unsqueeze(-1))
            if i != seq_len:
                next_token_emb = self.embedding(next_token)  # B * H
                transformer_input[:, i + 1] = next_token_emb

        all_tokens = torch.cat(all_tokens, dim=-1)  # B * S
        print(g.y[:2])
        print(all_tokens[:2])
        print("*" * 89)

        return all_tokens  # B * S


class TransformerModel(nn.Module):
    """Transformer encoder model using Pytorch.
    Args:
        input_dim (int): Input dimension of the model.
        num_layers (int): Number of transformer layer.
        hidden_dim (int): Hidden dimension in transformer model.
        num_heads (int): Number of head in each transformer layer.

    """

    def __init__(
        self, input_dim: int, num_layers: int, hidden_dim: int, num_heads: int
    ):
        super(TransformerModel, self).__init__()
        self.encoder = nn.TransformerEncoder(
            nn.TransformerEncoderLayer(
                input_dim, num_heads, hidden_dim, batch_first=True
            ),
            num_layers,
        )

    def forward(
        self,
        x: Tensor,
        mask: Tensor = None,
        src_key_padding_mask: Tensor = None,
    ) -> Tensor:
        encoded = self.encoder(
            x, mask=mask, src_key_padding_mask=src_key_padding_mask
        )
        return encoded


class PyGRGCNEdge(MultiLayerMessagePassing):
    def __init__(
        self,
        num_layers: int,
        num_rels: int,
        inp_dim: int,
        out_dim: int,
        drop_ratio=0,
        JK="last",
        batch_norm=True,
    ):
        super().__init__(
            num_layers, inp_dim, out_dim, drop_ratio, JK, batch_norm
        )
        self.num_rels = num_rels
        self.build_layers()

    def build_input_layer(self):
        return RGCNEdgeConv(self.inp_dim, self.out_dim, self.num_rels)

    def build_hidden_layer(self):
        return RGCNEdgeConv(self.inp_dim, self.out_dim, self.num_rels)

    def build_message_from_input(self, g):
        return {
            "g": g.edge_index,
            "h": g.x,
            "e": g.edge_type,
            "he": g.edge_attr,
        }

    def build_message_from_output(self, g, h):
        return {"g": g.edge_index, "h": h, "e": g.edge_type, "he": g.edge_attr}

    def layer_forward(self, layer, message):
        return self.conv[layer](
            message["h"], message["he"], message["g"], message["e"]
        )

models/model.py
LechengKong/OneForAll
from gp.utils.datasets import DatasetWithCollate
import torch_geometric as pyg
import torch
import numpy as np
from scipy.sparse import csr_array, coo_array
from gp.utils.graph import sample_fixed_hop_size_neighbor
from fs_datamanager import FewShotDataManager
from torch_geometric.data import Batch
from torch_geometric.utils import dropout_edge
from utils import scipy_rwpe

from gp.utils.utils import SmartTimer


class GraphTextDataset(DatasetWithCollate):
    def __init__(self, graph, process_label_func, **kwargs):
        self.g = graph
        self.process_label_func = process_label_func
        self.kwargs = kwargs

    def __getitem__(self, index):
        feature_graph = self.make_feature_graph(index)
        prompt_graph = self.make_prompted_graph(feature_graph)
        ret_data = self.to_pyg(feature_graph, prompt_graph)
        if (
            "walk_length" in self.kwargs
            and self.kwargs["walk_length"] is not None
        ):
            ret_data.rwpe = scipy_rwpe(ret_data, self.kwargs["walk_length"])
        return ret_data

    def make_feature_graph(self, index):
        pass

    def make_prompted_graph(self, feature_graph):
        pass

    def to_pyg(self, feature_graph, prompted_graph):
        pass

    def get_collate_fn(self):
        return pyg.loader.dataloader.Collater(None, None)

    def process_label(self, label):
        if self.process_label_func is None:
            trimed_class = torch.zeros((1, len(self.class_emb)))
            trimed_class[0, label] = 1
            return label, self.class_emb, trimed_class
        else:
            return self.process_label_func(self.class_emb, label)


class SubgraphDataset(GraphTextDataset):
    def __init__(
        self,
        pyg_graph,
        class_emb,
        data_idx,
        hop=2,
        class_mapping=None,
        to_undirected=False,
        process_label_func=None,
        adj=None,
        **kwargs,
    ):
        super().__init__(pyg_graph, process_label_func, **kwargs)
        self.to_undirected = to_undirected
        edge_index = self.g.edge_index
        if self.to_undirected:
            edge_index = pyg.utils.to_undirected(edge_index)
        if adj is not None:
            self.adj = adj
        else:
            self.adj = csr_array(
                (
                    torch.ones(len(edge_index[0])),
                    (edge_index[0], edge_index[1]),
                ),
                shape=(self.g.num_nodes, self.g.num_nodes),
            )
        self.class_emb = class_emb
        self.hop = hop
        self.data_idx = data_idx
        self.class_mapping = class_mapping

    def __len__(self):
        return len(self.data_idx)

    def get_neighbors(self, index):
        node_id = self.data_idx[index]
        neighbors = sample_fixed_hop_size_neighbor(
            self.adj, [node_id], self.hop, max_nodes_per_hope=100
        )
        neighbors = np.r_[node_id, neighbors]
        edges = self.adj[neighbors, :][:, neighbors].tocoo()
        if self.class_mapping is not None:
            label = self.class_mapping[self.g.y[node_id]]
        else:
            label = self.g.y[node_id]
        edge_index = torch.stack(
            [
                torch.tensor(edges.row, dtype=torch.long),
                torch.tensor(edges.col, dtype=torch.long),
            ]
        )
        label, emb, binary_rep = self.process_label(label)
        return edge_index, neighbors, label, binary_rep, [0]

    def make_feature_graph(self, index):
        (
            edge_index,
            neighbors,
            label,
            binary_rep,
            target_node_id,
        ) = self.get_neighbors(index)
        feat = self.g.x_text_feat[neighbors]
        e_type = torch.zeros(len(edge_index[0]), dtype=torch.long)
        edge_feat = self.g.edge_text_feat.repeat([len(edge_index[0]), 1])
        return (
            feat,
            edge_feat,
            edge_index,
            e_type,
            target_node_id,
            label,
            binary_rep,
        )

    def make_prompted_graph(self, feature_graph):
        (
            feat,
            edge_feat,
            edge_index,
            e_type,
            target_node_id,
            label,
            binary_rep,
        ) = feature_graph
        next_nid = len(feat)
        feat = torch.cat([feat, self.class_emb], dim=0)
        virtual_edge = torch.tensor(
            [
                target_node_id * len(self.class_emb),
                [i + next_nid for i in range(len(self.class_emb))],
            ],
            dtype=torch.long,
        )
        edge_index = torch.cat(
            [edge_index, virtual_edge, virtual_edge[[1, 0]]],
            dim=-1,
        )
        e_type = torch.cat(
            [
                e_type,
                torch.zeros(len(virtual_edge[0]), dtype=torch.long) + 1,
                torch.zeros(len(virtual_edge[0]), dtype=torch.long) + 2,
            ]
        )
        edge_feat = torch.cat(
            [
                edge_feat,
                self.g.prompt_edge_feat.repeat([len(virtual_edge[0]) * 2, 1]),
            ]
        )
        new_subg = pyg.data.Data(
            feat, edge_index, y=label, edge_attr=edge_feat, edge_type=e_type
        )
        return new_subg

    def to_pyg(self, feature_graph, prompted_graph):
        prompt_nodes_mask = torch.zeros(
            prompted_graph.num_nodes, dtype=torch.bool
        )
        bin_labels = torch.zeros(prompted_graph.num_nodes, dtype=torch.float)
        prompt_nodes_mask[
            prompted_graph.num_nodes - len(self.class_emb) :
        ] = True
        # noi_node_mask = torch.zeros(prompted_graph.num_nodes, dtype=torch.bool)
        # noi_node_mask[
        #     prompted_graph.num_nodes - len(self.class_emb) - 1
        # ] = True
        # prompted_graph.noi_node_mask = noi_node_mask
        prompted_graph.true_nodes_mask = prompt_nodes_mask
        bin_labels[
            prompted_graph.num_nodes - len(self.class_emb) :
        ] = feature_graph[-1]
        prompted_graph.bin_labels = bin_labels
        target_node_mask = torch.zeros(
            prompted_graph.num_nodes, dtype=torch.bool
        )
        target_node_mask[feature_graph[-3]] = True
        prompted_graph.target_node_mask = target_node_mask
        prompted_graph.sample_num_nodes = prompted_graph.num_nodes
        prompted_graph.num_classes = len(self.class_emb)
        # print("text", new_subg)
        return prompted_graph


class SubgraphNopromptDataset(SubgraphDataset):
    def make_prompted_graph(self, feature_graph):
        (
            feat,
            edge_feat,
            edge_index,
            e_type,
            target_node_id,
            label,
            binary_rep,
        ) = feature_graph
        feat = torch.cat([feat, self.class_emb], dim=0)
        new_subg = pyg.data.Data(
            feat, edge_index, y=label, edge_attr=edge_feat, edge_type=e_type
        )
        return new_subg


class SubgraphHierDataset(SubgraphDataset):
    def __init__(
        self,
        pyg_graph,
        class_emb,
        data_idx,
        hop=2,
        class_mapping=None,
        prompt_feat=None,
        to_undirected=False,
        process_label_func=None,
        adj=None,
        **kwargs,
    ):
        super().__init__(
            pyg_graph,
            class_emb,
            data_idx,
            hop,
            class_mapping,
            to_undirected,
            process_label_func,
            adj,
            **kwargs,
        )
        self.prompt_feat = prompt_feat

    def __len__(self):
        return len(self.data_idx)

    def make_prompted_graph(self, feature_graph):
        (
            feat,
            edge_feat,
            edge_index,
            e_type,
            target_node_id,
            label,
            binary_rep,
        ) = feature_graph
        next_nid = len(feat)
        feat = torch.cat([feat, self.prompt_feat, self.class_emb], dim=0)
        virtual_edge = torch.tensor(
            [
                target_node_id + [next_nid] * len(self.class_emb),
                [next_nid] * len(target_node_id)
                + [i + next_nid + 1 for i in range(len(self.class_emb))],
            ],
            dtype=torch.long,
        )
        edge_index = torch.cat(
            [edge_index, virtual_edge, virtual_edge[[1, 0]]],
            dim=-1,
        )
        e_type = torch.cat(
            [
                e_type,
                torch.zeros(1, dtype=torch.long) + 1,
                torch.zeros(int(len(virtual_edge[0]) - 1), dtype=torch.long)
                + 2,
                torch.zeros(1, dtype=torch.long) + 3,
                torch.zeros(int(len(virtual_edge[0]) - 1), dtype=torch.long)
                + 4,
            ]
        )
        edge_feat = torch.cat(
            [
                edge_feat,
                self.g.prompt_edge_feat.repeat([len(virtual_edge[0]) * 2, 1]),
            ]
        )
        # edge_index = torch.cat([edge_index, edge_index[[1, 0]]], dim=-1)
        new_subg = pyg.data.Data(
            feat, edge_index, y=label, edge_attr=edge_feat, edge_type=e_type
        )
        return new_subg


class SubgraphLinkHierDataset(SubgraphHierDataset):
    def __init__(
        self,
        pyg_graph,
        class_emb,
        edges,
        remove_edge=False,
        hop=2,
        class_mapping=None,
        prompt_feat=None,
        to_undirected=False,
        process_label_func=None,
        adj=None,
        **kwargs,
    ):
        super().__init__(
            pyg_graph,
            class_emb,
            None,
            hop,
            class_mapping,
            prompt_feat,
            to_undirected,
            process_label_func,
            adj,
            **kwargs,
        )
        self.edges = edges
        self.pos_index = len(self.edges)
        self.remove_edge = remove_edge
        dense_adj = self.adj.todense() == 0
        neg_row, neg_col = np.nonzero(dense_adj)
        neg_edge_idx = np.random.permutation(len(neg_row))[: self.pos_index]
        neg_row, neg_col = neg_row[neg_edge_idx], neg_col[neg_edge_idx]
        self.neg_edges = np.stack([neg_row, neg_col], axis=1)

        self.total_edges = np.concatenate([self.edges, self.neg_edges], axis=0)

    def __len__(self):
        return len(self.total_edges)

    def remove_link(self, row, col):
        remove_ind = np.logical_or(
            np.logical_and(row == 0, col == 1),
            np.logical_and(row == 1, col == 0),
        )
        keep_ind = np.logical_not(remove_ind)
        return row[keep_ind], col[keep_ind]

    def get_neighbors(self, index):
        edge_id = self.total_edges[index]

        if index < self.pos_index:
            label = 1
        else:
            label = 0
        node_ids = list(edge_id)
        neighbors = sample_fixed_hop_size_neighbor(
            self.adj, node_ids, self.hop, max_nodes_per_hope=100
        )
        neighbors = np.r_[node_ids, neighbors]
        edges = self.adj[neighbors, :][:, neighbors].tocoo()
        row = edges.row
        col = edges.col
        if self.remove_edge and index < self.pos_index:
            row, col = self.remove_link(row, col)
        edge_index = torch.stack(
            [
                torch.tensor(row, dtype=torch.long),
                torch.tensor(col, dtype=torch.long),
            ]
        )
        label, embs, binary_rep = self.process_label(label)
        return edge_index, neighbors, label, binary_rep, [0, 1]


class SubgraphKGHierDataset(SubgraphHierDataset):
    def __init__(
        self,
        pyg_graph,
        class_emb,
        edges,
        remove_edge=False,
        hop=2,
        class_mapping=None,
        prompt_feat=None,
        to_undirected=False,
        process_label_func=None,
        adj=None,
        **kwargs,
    ):
        super().__init__(
            pyg_graph,
            class_emb,
            None,
            hop,
            class_mapping,
            prompt_feat,
            to_undirected,
            process_label_func,
            adj,
            **kwargs,
        )
        self.edges = edges
        self.remove_edge = remove_edge

    def __len__(self):
        return len(self.edges[0])

    def index_to_mask(self, index, size=None):
        size = int(index.max()) + 1 if size is None else size
        mask = torch.zeros(size, dtype=torch.bool)
        mask[index] = True
        return mask

    def remove_link(self, row, col, val, target_idx):
        keep_ind = val != target_idx
        return row[keep_ind], col[keep_ind], val[keep_ind]

    def get_neighbors(self, index):
        node_ids = list(self.edges[0][index])
        label = self.edges[1][index]

        neighbors = sample_fixed_hop_size_neighbor(
            self.adj, node_ids, self.hop, max_nodes_per_hope=100
        )
        neighbors = np.r_[node_ids, neighbors]
        node_mask = self.index_to_mask(neighbors, size=self.g.num_nodes)

        edge_mask = (
            node_mask[self.g.edge_index[0]] & node_mask[self.g.edge_index[1]]
        )
        if self.remove_edge:
            index_mask = torch.ones(
                len(self.g.edge_index[0]), dtype=torch.bool
            )
            index_mask[index] = False
            edge_mask = edge_mask & index_mask
        edge2idx = torch.zeros(self.g.num_nodes, dtype=torch.long)
        edge2idx[neighbors] = torch.arange(len(neighbors))
        edge_index = self.g.edge_index[:, edge_mask]
        edge_type = self.g.edge_types[edge_mask]
        edge_index = edge2idx[edge_index]
        label, embs, binary_rep = self.process_label(label)
        return edge_index, neighbors, label, binary_rep, [0, 1], edge_type

    def make_feature_graph(self, index):
        (
            edge_index,
            neighbors,
            label,
            binary_rep,
            target_node_id,
            edge_type,
        ) = self.get_neighbors(index)
        edge_index = torch.cat([edge_index, edge_index[[1, 0]]], dim=-1)
        feat = self.g.x_text_feat[neighbors]
        e_type = torch.zeros(len(edge_index[0]), dtype=torch.long)
        edge_feat = self.g.edge_text_feat[
            torch.cat(
                [edge_type, edge_type + int(len(self.g.edge_text_feat) / 2)]
            )
        ]
        return (
            feat,
            edge_feat,
            edge_index,
            e_type,
            target_node_id,
            label,
            binary_rep,
        )


class GraphListDataset(GraphTextDataset):
    def __init__(
        self,
        graphs,
        class_embs,
        prompt_edge_feat,
        data_idx,
        process_label_func=None,
        single_prompt_edge=False,
        **kwargs,
    ):
        super().__init__(graphs, process_label_func, **kwargs)
        self.class_embs = class_embs
        self.prompt_edge_feat = prompt_edge_feat
        self.data_idx = data_idx
        self.single_prompt_edge = single_prompt_edge

    def __len__(self):
        return len(self.data_idx)

    def make_feature_graph(self, index):
        g = self.g[self.data_idx[index]]
        edge_index = g.edge_index
        label = g.y
        # label_emb = self.class_emb(label).view(1, -1)
        feat = g.x_text_feat
        next_nid = g.num_nodes
        edge_feat = g.xe_text_feat
        y_label, g_class_emb, trimmed_label = self.process_label(label)
        return (
            feat,
            edge_feat,
            edge_index,
            next_nid,
            g_class_emb,
            y_label,
            trimmed_label,
        )

    def make_prompted_graph(self, feature_graph):
        (
            feat,
            edge_feat,
            edge_index,
            next_nid,
            g_class_emb,
            label,
            trimmed_label,
        ) = feature_graph
        feat = torch.cat([feat, g_class_emb], dim=0)
        virtual_edge = torch.stack(
            [
                torch.arange(next_nid, dtype=torch.long)
                .repeat(1, len(g_class_emb))
                .view(-1),
                torch.arange(
                    next_nid, next_nid + len(g_class_emb), dtype=torch.long
                ).repeat_interleave(next_nid),
            ],
            dim=0,
        )
        # print(virtual_edge)
        edge_index = torch.cat(
            [
                edge_index,
                virtual_edge,
                # virtual_edge[[1, 0]],
            ],
            dim=-1,
        )
        edge_type = torch.cat(
            [
                torch.zeros(len(edge_feat), dtype=torch.long),
                torch.zeros(len(virtual_edge[0]), dtype=torch.long) + 1,
                # torch.zeros(len(virtual_edge[0]), dtype=torch.long) + 2,
            ]
        )
        edge_emb = torch.cat(
            [
                edge_feat,
                self.prompt_edge_feat[0].repeat(
                    [(next_nid) * len(g_class_emb), 1]
                ),
            ]
        )
        prompted_graph = pyg.data.Data(
            feat,
            edge_index,
            y=label,
            edge_attr=edge_emb,
            edge_type=edge_type,
        )
        # print(prompted_graph)
        # print(edge_index)
        return prompted_graph

    def to_pyg(self, feature_graph, prompted_graph):
        true_nodes_mask = torch.zeros(
            prompted_graph.num_nodes, dtype=torch.bool
        )
        true_nodes_mask[
            prompted_graph.num_nodes - len(feature_graph[-3]) :
        ] = True
        bin_labels = torch.zeros(prompted_graph.num_nodes, dtype=torch.float)
        bin_labels[
            prompted_graph.num_nodes - len(feature_graph[-3]) :
        ] = feature_graph[-1]
        # noi_node_mask = torch.zeros(prompted_graph.num_nodes, dtype=torch.bool)
        # noi_node_mask[
        #     prompted_graph.num_nodes - len(feature_graph[-3]) - 1
        # ] = True
        # prompted_graph.noi_node_mask = noi_node_mask
        prompted_graph.bin_labels = bin_labels
        target_node_mask = torch.zeros(
            prompted_graph.num_nodes, dtype=torch.bool
        )
        target_node_mask[: feature_graph[-4]] = True
        prompted_graph.target_node_mask = target_node_mask
        prompted_graph.true_nodes_mask = true_nodes_mask
        prompted_graph.sample_num_nodes = prompted_graph.num_nodes
        prompted_graph.num_classes = len(feature_graph[-3])
        return prompted_graph


class GraphListNopromptDataset(GraphListDataset):
    def make_prompted_graph(self, feature_graph):
        (
            feat,
            edge_feat,
            edge_index,
            next_nid,
            g_class_emb,
            label,
            trimmed_label,
        ) = feature_graph
        feat = torch.cat([feat, g_class_emb], dim=0)
        edge_type = torch.zeros(len(edge_feat), dtype=torch.long)
        prompted_graph = pyg.data.Data(
            feat,
            edge_index,
            y=label,
            edge_attr=edge_feat,
            edge_type=edge_type,
        )
        # print(prompted_graph)
        # print(edge_index)
        return prompted_graph


class GraphListHierDataset(GraphListDataset):
    def __init__(
        self,
        graphs,
        class_embs,
        prompt_edge_feat,
        prompt_text_feat,
        data_idx,
        process_label_func=None,
        single_prompt_edge=False,
        **kwargs,
    ):
        super().__init__(
            graphs,
            class_embs,
            prompt_edge_feat,
            data_idx,
            process_label_func,
            single_prompt_edge,
            **kwargs,
        )
        self.prompt_text_feat = prompt_text_feat

    def make_prompted_graph(self, feature_graph):
        (
            feat,
            edge_feat,
            edge_index,
            next_nid,
            g_class_emb,
            label,
            trimmed_label,
        ) = feature_graph
        feat = torch.cat([feat, self.prompt_text_feat[:1], g_class_emb], dim=0)
        feature2prompt_edge = torch.tensor(
            [list(range(next_nid)), [next_nid] * next_nid],
            dtype=torch.long,
        )
        prompt2prompt_edge = torch.tensor(
            [
                [next_nid] * len(g_class_emb),
                [next_nid + i + 1 for i in range(len(g_class_emb))],
            ],
            dtype=torch.long,
        )
        prompt_edges = [
            feature2prompt_edge,
            feature2prompt_edge[[1, 0]],
            prompt2prompt_edge,
        ]
        prompt_edge_types = [
            torch.zeros(
                len(feature2prompt_edge[0]),
                dtype=torch.long,
            )
            + 1,
            torch.zeros(len(feature2prompt_edge[0]), dtype=torch.long) + 2,
            torch.zeros(len(prompt2prompt_edge[0]), dtype=torch.long) + 3,
        ]
        prompt_feat_multiple = 1
        if not self.single_prompt_edge:
            prompt_edges += [prompt2prompt_edge[[1, 0]]]
            prompt_edge_types += [
                torch.zeros(len(prompt2prompt_edge[0]), dtype=torch.long) + 4
            ]
            prompt_feat_multiple = 2
        edge_index = torch.cat([edge_index] + prompt_edges, dim=-1)
        edge_type = torch.cat(
            [torch.zeros(len(edge_feat), dtype=torch.long)] + prompt_edge_types
        )
        edge_feat = torch.cat(
            [
                edge_feat,
                self.prompt_edge_feat[0].repeat(
                    [
                        2 * len(feature2prompt_edge[0])
                        + len(prompt2prompt_edge[0]) * prompt_feat_multiple,
                        1,
                    ]
                ),
            ]
        )
        prompted_graph = pyg.data.Data(
            feat,
            edge_index,
            y=label,
            edge_attr=edge_feat,
            edge_type=edge_type,
        )
        return prompted_graph


class GraphListHierFSDataset(GraphListDataset):
    def __init__(
        self,
        graphs,
        class_embs,
        prompt_edge_feat,
        prompt_text_feat,
        data_idx,
        process_label_func=None,
        single_prompt_edge=False,
        class_ind=None,
        max_sample=10,
        shot=None,
        reuse=False,
        use_class_emb=False,
        target_class=None,
        **kwargs,
    ):
        super().__init__(
            graphs,
            class_embs,
            prompt_edge_feat,
            data_idx,
            process_label_func,
            single_prompt_edge,
            **kwargs,
        )
        self.prompt_text_feat = prompt_text_feat
        self.class_ind = class_ind
        self.max_sample = max_sample
        self.num_classes = self.class_ind.size()[1]
        self.shot = shot
        self.reuse = reuse
        self.true_can = [
            (can == 1).nonzero(as_tuple=True)[0] for can in self.class_ind.T
        ]
        self.false_can = [
            (can == 0).nonzero(as_tuple=True)[0] for can in self.class_ind.T
        ]
        self.use_class_emb = use_class_emb
        self.target_class = target_class

    def make_feature_graph(self, index):
        if self.target_class is not None:
            class_ind = self.target_class[
                torch.randint(0, len(self.target_class), (1,))
            ]
        else:
            class_ind = torch.randint(0, self.num_classes, (1,))
        if self.shot is None:
            while (
                len(self.true_can[class_ind]) <= 1
                or len(self.false_can[class_ind]) <= 1
            ):
                class_ind = torch.randint(0, self.num_classes, (1,))
            c_max_sample = min(
                len(self.true_can[class_ind]), len((self.false_can[class_ind]))
            )
            num_sample = torch.randint(
                2, min(c_max_sample, self.max_sample) + 1, (1,)
            )
        else:
            while (
                len(self.true_can[class_ind]) <= self.shot
                or len(self.false_can[class_ind]) <= self.shot
            ):
                class_ind = torch.randint(0, self.num_classes, (1,))
            num_sample = self.shot + 1
        # print(num_sample)
        true_can = self.true_can[class_ind][
            torch.randperm(len(self.true_can[class_ind]))[:num_sample]
        ]
        # print(true_can)
        false_can = self.false_can[class_ind][
            torch.randperm(len(self.false_can[class_ind]))[:num_sample]
        ]
        true_feature_graph = [
            super(GraphListHierFSDataset, self).make_feature_graph(can)
            for can in true_can
        ]
        false_feature_graph = [
            super(GraphListHierFSDataset, self).make_feature_graph(can)
            for can in false_can
        ]
        return (
            true_feature_graph,
            false_feature_graph,
            self.class_embs[class_ind],
        )

    def self_sample_edge_index(self, init_ind, size):
        prompt_init_ind = init_ind + size * 2
        self_edges = torch.stack(
            [
                torch.arange(init_ind, prompt_init_ind),
                torch.arange(prompt_init_ind, prompt_init_ind + size * 2),
            ],
            dim=0,
        )
        fs_edges = []
        for i in range(size):
            fs_edges.append(
                torch.stack(
                    [
                        i
                        + init_ind
                        + torch.zeros(2 * size - 2, dtype=torch.long),
                        torch.cat(
                            [
                                torch.arange(
                                    prompt_init_ind,
                                    i + prompt_init_ind,
                                    dtype=torch.long,
                                ),
                                torch.arange(
                                    i + prompt_init_ind + 1,
                                    prompt_init_ind + size,
                                    dtype=torch.long,
                                ),
                                torch.arange(
                                    prompt_init_ind + size,
                                    prompt_init_ind + size + i,
                                    dtype=torch.long,
                                ),
                                torch.arange(
                                    prompt_init_ind + size + i + 1,
                                    prompt_init_ind + 2 * size,
                                    dtype=torch.long,
                                ),
                            ]
                        ),
                    ],
                    dim=0,
                )
            )
        fs_edges = torch.cat(fs_edges, dim=-1)
        return self_edges, fs_edges

    def make_prompted_graph(self, feature_graph):
        if self.reuse:
            return self.make_prompted_graph_reuse(feature_graph)
        else:
            return self.make_prompted_graph_simple(feature_graph)

    def make_prompted_graph_simple(self, feature_graph):
        (true_feature_graph, false_feature_graph, class_emb) = feature_graph
        if torch.rand(1)[0] < 0.5:
            target_graph = true_feature_graph[0]
            sample_label = torch.tensor([[1]])
        else:
            target_graph = false_feature_graph[0]
            sample_label = torch.tensor([[0]])

        all_graphs = [target_graph] + true_feature_graph[1:]
        num_nodes = torch.tensor(
            [len(g[0]) for g in all_graphs], dtype=torch.long
        )
        total_num_nodes = num_nodes.sum()
        feat = torch.cat(
            [g[0] for g in all_graphs]
            + [
                self.prompt_text_feat[:1].repeat(len(all_graphs), 1),
                # class_emb,
                class_emb
                if self.use_class_emb
                else self.prompt_text_feat[1:2].repeat(len(class_emb), 1),
            ],
            dim=0,
        )

        feature_nodes_indices = torch.arange(total_num_nodes, dtype=torch.long)
        prompt_nodes_indices = torch.arange(
            total_num_nodes,
            total_num_nodes + len(all_graphs),
            dtype=torch.long,
        ).repeat_interleave(num_nodes)

        feature2prompt_edge = torch.stack(
            [feature_nodes_indices, prompt_nodes_indices], dim=0
        )
        valid_label_ind = total_num_nodes + len(all_graphs)
        query_prompt_edge = torch.stack(
            [
                torch.zeros(1, dtype=torch.long) + total_num_nodes,
                torch.arange(
                    total_num_nodes + len(all_graphs),
                    total_num_nodes + len(all_graphs) + 1,
                    dtype=torch.long,
                ),
            ],
            dim=0,
        )
        support_prompt_edge = torch.stack(
            [
                torch.arange(
                    total_num_nodes + 1,
                    total_num_nodes + len(all_graphs),
                    dtype=torch.long,
                ),
                valid_label_ind.repeat_interleave(len(all_graphs) - 1),
            ],
            dim=0,
        )

        prompt2prompt_edge = torch.cat(
            [query_prompt_edge, support_prompt_edge],
            dim=-1,
        )
        prompt_edges = [
            feature2prompt_edge,
            feature2prompt_edge[[1, 0]],
            prompt2prompt_edge,
        ]
        prompt_edge_types = [
            torch.zeros(
                len(feature2prompt_edge[0]),
                dtype=torch.long,
            )
            + 1,
            torch.zeros(len(feature2prompt_edge[0]), dtype=torch.long) + 2,
            torch.zeros(len(query_prompt_edge[0]), dtype=torch.long) + 3,
            torch.zeros(len(support_prompt_edge[0]), dtype=torch.long) + 4,
        ]
        prompt_feat_multiple = 1
        if not self.single_prompt_edge:
            prompt_edges += [prompt2prompt_edge[[1, 0]]]
            prompt_edge_types += [
                torch.zeros(len(prompt2prompt_edge[0]), dtype=torch.long) + 4
            ]
            prompt_feat_multiple = 2
        edge_feat = torch.cat([g[1] for g in all_graphs], dim=0)
        edge_index = torch.cat([g[2] for g in all_graphs], dim=-1)
        edge_index = torch.cat([edge_index] + prompt_edges, dim=-1)
        edge_type = torch.cat(
            [torch.zeros(len(edge_feat), dtype=torch.long)] + prompt_edge_types
        )
        edge_feat = torch.cat(
            [
                edge_feat,
                self.prompt_edge_feat[0].repeat(
                    [
                        2 * len(feature2prompt_edge[0]),
                        1,
                    ]
                ),
                self.prompt_edge_feat[1].repeat(
                    [
                        len(query_prompt_edge[0]),
                        1,
                    ]
                ),
                self.prompt_edge_feat[2].repeat(
                    [
                        len(support_prompt_edge[0]),
                        1,
                    ]
                ),
            ]
        )
        prompted_graph = pyg.data.Data(
            feat,
            edge_index,
            y=sample_label,
            edge_attr=edge_feat,
            edge_type=edge_type,
        )
        return prompted_graph

    def make_prompted_graph_reuse(self, feature_graph):
        (true_feature_graph, false_feature_graph, class_emb) = feature_graph
        num_sample = len(false_feature_graph)
        all_graphs = true_feature_graph + false_feature_graph
        num_nodes = torch.tensor(
            [len(g[0]) for g in all_graphs], dtype=torch.long
        )
        total_num_nodes = num_nodes.sum()
        # print(self.prompt_text_feat)
        feat = torch.cat(
            [g[0] for g in all_graphs]
            + [
                self.prompt_text_feat.repeat(len(all_graphs), 1),
                class_emb.repeat(len(all_graphs), 1),
            ],
            dim=0,
        )
        feature_nodes_indices = torch.arange(total_num_nodes, dtype=torch.long)
        prompt_nodes_indices = torch.arange(
            total_num_nodes,
            total_num_nodes + len(all_graphs),
            dtype=torch.long,
        ).repeat_interleave(num_nodes)
        feature2prompt_edge = torch.stack(
            [feature_nodes_indices, prompt_nodes_indices], dim=0
        )
        (
            self_prompt2prompt_edge,
            fs_prompt2prompt_edge,
        ) = self.self_sample_edge_index(total_num_nodes, num_sample)
        prompt2prompt_edge = torch.cat(
            [self_prompt2prompt_edge, fs_prompt2prompt_edge],
            dim=-1,
        )
        prompt_edges = [
            feature2prompt_edge,
            feature2prompt_edge[[1, 0]],
            prompt2prompt_edge,
        ]
        prompt_edge_types = [
            torch.zeros(
                len(feature2prompt_edge[0]),
                dtype=torch.long,
            )
            + 1,
            torch.zeros(len(feature2prompt_edge[0]), dtype=torch.long) + 2,
            torch.zeros(len(prompt2prompt_edge[0]), dtype=torch.long) + 3,
        ]
        prompt_feat_multiple = 1
        if not self.single_prompt_edge:
            prompt_edges += [prompt2prompt_edge[[1, 0]]]
            prompt_edge_types += [
                torch.zeros(len(prompt2prompt_edge[0]), dtype=torch.long) + 4
            ]
            prompt_feat_multiple = 2
        edge_feat = torch.cat([g[1] for g in all_graphs], dim=0)
        edge_index = torch.cat([g[2] for g in all_graphs], dim=-1)
        edge_index = torch.cat([edge_index] + prompt_edges, dim=-1)
        edge_type = torch.cat(
            [torch.zeros(len(edge_feat), dtype=torch.long)] + prompt_edge_types
        )
        edge_feat = torch.cat(
            [
                edge_feat,
                self.prompt_edge_feat[0].repeat(
                    [
                        2 * len(feature2prompt_edge[0]),
                        1,
                    ]
                ),
                self.prompt_edge_feat[1].repeat(
                    [
                        len(self_prompt2prompt_edge[0]),
                        1,
                    ]
                ),
                self.prompt_edge_feat[2].repeat(
                    [
                        len(fs_prompt2prompt_edge[0]),
                        1,
                    ]
                ),
            ]
        )
        prompted_graph = pyg.data.Data(
            feat,
            edge_index,
            y=torch.cat([g[-2] for g in all_graphs], dim=0),
            edge_attr=edge_feat,
            edge_type=edge_type,
        )
        return prompted_graph

    def to_pyg(self, feature_graph, prompted_graph):
        if self.reuse:
            return self.to_pyg_reuse(feature_graph, prompted_graph)
        else:
            return self.to_pyg_simple(feature_graph, prompted_graph)

    def to_pyg_reuse(self, feature_graph, prompted_graph):
        true_nodes_mask = torch.zeros(
            prompted_graph.num_nodes, dtype=torch.bool
        )
        true_nodes_mask[
            prompted_graph.num_nodes - len(feature_graph[1]) * 2 :
        ] = True
        bin_labels = torch.zeros(prompted_graph.num_nodes, dtype=torch.float)
        bin_labels[
            prompted_graph.num_nodes
            - len(feature_graph[1]) * 2 : prompted_graph.num_nodes
            - len(feature_graph[1])
        ] = 1
        # bin_labels[prompted_graph.num_nodes - len(feature_graph[1]) :] = 0
        prompted_graph.bin_labels = bin_labels
        target_node_mask = torch.zeros(
            prompted_graph.num_nodes, dtype=torch.bool
        )
        # target_node_mask[feature_graph[0][-4]] = True
        prompted_graph.target_node_mask = target_node_mask
        prompted_graph.true_nodes_mask = true_nodes_mask
        prompted_graph.sample_num_nodes = prompted_graph.num_nodes
        prompted_graph.num_classes = 1
        return prompted_graph

    def to_pyg_simple(self, feature_graph, prompted_graph):
        true_nodes_mask = torch.zeros(
            prompted_graph.num_nodes, dtype=torch.bool
        )
        true_nodes_mask[prompted_graph.num_nodes - 1 :] = True
        bin_labels = torch.zeros(prompted_graph.num_nodes, dtype=torch.float)
        bin_labels[prompted_graph.num_nodes - 1 :] = prompted_graph.y
        # bin_labels[prompted_graph.num_nodes - len(feature_graph[1]) :] = 0
        prompted_graph.bin_labels = bin_labels
        target_node_mask = torch.zeros(
            prompted_graph.num_nodes, dtype=torch.bool
        )
        # target_node_mask[feature_graph[0][-4]] = True
        prompted_graph.target_node_mask = target_node_mask
        prompted_graph.true_nodes_mask = true_nodes_mask
        prompted_graph.sample_num_nodes = prompted_graph.num_nodes
        prompted_graph.num_classes = 1
        return prompted_graph


class GraphListHierFixDataset(GraphListDataset):
    def __init__(
        self,
        graphs,
        class_embs,
        prompt_edge_feat,
        prompt_text_feat,
        data_idx,
        process_label_func=None,
        single_prompt_edge=False,
        class_ind=None,
        shot=10,
        use_class_emb=False,
        **kwargs,
    ):
        super().__init__(
            graphs,
            class_embs,
            prompt_edge_feat,
            data_idx,
            process_label_func,
            single_prompt_edge,
            **kwargs,
        )
        self.prompt_text_feat = prompt_text_feat
        self.class_ind = class_ind
        self.shot = shot
        self.num_classes = self.class_ind.size()[1]
        self.true_can = [
            (can == 1).nonzero(as_tuple=True)[0] for can in self.class_ind.T
        ]
        self.false_can = [
            (can == 0).nonzero(as_tuple=True)[0] for can in self.class_ind.T
        ]
        self.use_class_emb = use_class_emb

    def make_feature_graph(self, index):
        query_graph = super(GraphListHierFixDataset, self).make_feature_graph(
            index
        )
        # print(true_can)
        label = query_graph[-1]
        support_graph = []
        valid_labels = []
        for i, label_val in enumerate(label.view(-1)):
            # print(i)
            if label_val == label_val:
                true_can = self.true_can[i]
                if len(true_can) < self.shot:
                    support_graph.append([])
                    label_val = float("nan")
                else:
                    support_g = true_can[
                        torch.randperm(len(true_can))[: self.shot]
                    ]
                    support_graph.append(
                        [
                            super(
                                GraphListHierFixDataset, self
                            ).make_feature_graph(g_ind)
                            for g_ind in support_g
                        ]
                    )
            valid_labels.append(label_val)
        true_label = torch.tensor(valid_labels, dtype=torch.float).view(1, -1)
        return (
            query_graph,
            support_graph,
            query_graph[-3],
            true_label,
        )

    def make_prompted_graph(self, feature_graph):
        (query_graph, support_graph, class_emb, true_label) = feature_graph
        all_graphs = [query_graph]
        for sg in support_graph:
            all_graphs += sg
        num_nodes = torch.tensor(
            [len(g[0]) for g in all_graphs], dtype=torch.long
        )
        total_num_nodes = num_nodes.sum()
        feat = torch.cat(
            [g[0] for g in all_graphs]
            + [
                self.prompt_text_feat[:1].repeat(len(all_graphs), 1),
                class_emb
                if self.use_class_emb
                else self.prompt_text_feat[1:2].repeat(len(class_emb), 1),
            ],
            dim=0,
        )
        feature_nodes_indices = torch.arange(total_num_nodes, dtype=torch.long)
        prompt_nodes_indices = torch.arange(
            total_num_nodes,
            total_num_nodes + len(all_graphs),
            dtype=torch.long,
        ).repeat_interleave(num_nodes)

        feature2prompt_edge = torch.stack(
            [feature_nodes_indices, prompt_nodes_indices], dim=0
        )
        valid_label_ind = (
            (true_label == true_label).nonzero(as_tuple=True)[1]
            + total_num_nodes
            + len(all_graphs)
        )
        query_prompt_edge = torch.stack(
            [
                torch.zeros(len(class_emb), dtype=torch.long)
                + total_num_nodes,
                torch.arange(
                    total_num_nodes + len(all_graphs),
                    total_num_nodes + len(all_graphs) + len(class_emb),
                    dtype=torch.long,
                ),
            ],
            dim=0,
        )
        support_prompt_edge = torch.stack(
            [
                torch.arange(
                    total_num_nodes + 1,
                    total_num_nodes + len(all_graphs),
                    dtype=torch.long,
                ),
                valid_label_ind.repeat_interleave(self.shot),
            ],
            dim=0,
        )

        prompt2prompt_edge = torch.cat(
            [query_prompt_edge, support_prompt_edge],
            dim=-1,
        )
        prompt_edges = [
            feature2prompt_edge,
            feature2prompt_edge[[1, 0]],
            prompt2prompt_edge,
        ]
        prompt_edge_types = [
            torch.zeros(
                len(feature2prompt_edge[0]),
                dtype=torch.long,
            )
            + 1,
            torch.zeros(len(feature2prompt_edge[0]), dtype=torch.long) + 2,
            torch.zeros(len(query_prompt_edge[0]), dtype=torch.long) + 3,
            torch.zeros(len(support_prompt_edge[0]), dtype=torch.long) + 3,
        ]
        prompt_feat_multiple = 1
        if not self.single_prompt_edge:
            prompt_edges += [prompt2prompt_edge[[1, 0]]]
            prompt_edge_types += [
                torch.zeros(len(prompt2prompt_edge[0]), dtype=torch.long) + 4
            ]
            prompt_feat_multiple = 2
        edge_feat = torch.cat([g[1] for g in all_graphs], dim=0)
        edge_index = torch.cat([g[2] for g in all_graphs], dim=-1)
        edge_index = torch.cat([edge_index] + prompt_edges, dim=-1)
        edge_type = torch.cat(
            [torch.zeros(len(edge_feat), dtype=torch.long)] + prompt_edge_types
        )
        edge_feat = torch.cat(
            [
                edge_feat,
                self.prompt_edge_feat[0].repeat(
                    [
                        2 * len(feature2prompt_edge[0]),
                        1,
                    ]
                ),
                self.prompt_edge_feat[1].repeat(
                    [
                        len(query_prompt_edge[0]),
                        1,
                    ]
                ),
                self.prompt_edge_feat[2].repeat(
                    [
                        len(support_prompt_edge[0]),
                        1,
                    ]
                ),
            ]
        )
        prompted_graph = pyg.data.Data(
            feat,
            edge_index,
            y=torch.cat([g[-2] for g in all_graphs], dim=0),
            edge_attr=edge_feat,
            edge_type=edge_type,
        )
        return prompted_graph

    def to_pyg(self, feature_graph, prompted_graph):
        true_nodes_mask = torch.zeros(
            prompted_graph.num_nodes, dtype=torch.bool
        )
        true_nodes_mask[
            prompted_graph.num_nodes - len(feature_graph[-2]) :
        ] = True
        bin_labels = torch.zeros(prompted_graph.num_nodes, dtype=torch.float)
        bin_labels[
            prompted_graph.num_nodes - len(feature_graph[-2]) :
        ] = feature_graph[-1]
        prompted_graph.bin_labels = bin_labels
        target_node_mask = torch.zeros(
            prompted_graph.num_nodes, dtype=torch.bool
        )
        # target_node_mask[feature_graph[0][-4]] = True
        prompted_graph.target_node_mask = target_node_mask
        prompted_graph.true_nodes_mask = true_nodes_mask
        prompted_graph.sample_num_nodes = prompted_graph.num_nodes
        prompted_graph.num_classes = len(feature_graph[-2])
        return prompted_graph


class FewShotSubgraphDataset(SubgraphDataset):
    def __init__(
        self,
        pyg_graph,
        class_emb,
        data_idx,
        n_way: int,
        k_shot: int,
        q_query: int,
        datamanager: FewShotDataManager,
        mode: int,
        hop=2,
        class_mapping=None,
        prompt_feat=None,
        to_undirected=False,
        adj=None,
        single_prompt_edge=False,
        **kwargs,
    ):
        super().__init__(
            pyg_graph,
            class_emb,
            data_idx,
            hop,
            class_mapping,
            to_undirected,
            adj=adj,
            **kwargs,
        )
        # mode 0 for sample index from training classes, 1 for val, 2 for test
        self.fs_idx_loader = datamanager.get_data_loader(mode)
        self.n_way = n_way
        self.k_shot = k_shot
        self.q_query = q_query
        self.prompt_feat = prompt_feat
        self.single_prompt_edge = single_prompt_edge

    def get_single_subgraph(self, node_id):
        """
        Get the info of node_id's khop subgraph
        Args:
            node_id: index of the node

        Returns:
            neighbors: feature matrix of all nodes in the subgraph
            edge_index: edge index tensor of the subgraph
            num_nodes: number of nodes in the subgraph
        """
        neighbors = sample_fixed_hop_size_neighbor(
            self.adj, [node_id], self.hop, max_nodes_per_hope=100
        )
        neighbors = np.r_[node_id, neighbors]
        edges = self.adj[neighbors, :][:, neighbors].tocoo()
        edge_index = torch.stack(
            [
                torch.tensor(edges.row, dtype=torch.long),
                torch.tensor(edges.col, dtype=torch.long),
            ]
        )
        return neighbors, edge_index, neighbors.shape[0]

    def make_feature_graph(self, graph_set):
        """
        For each query node, combine its subgraph with all support subgraphs.
        Args:
            qry_subgraph:
            spt_subgraphs:

        Returns:

        """
        (qry_subgraph, spt_subgraphs) = graph_set
        neighbors, edge_index, num_nodes = qry_subgraph
        neighbors_list = [neighbors] + [
            subgraph[0] for subgraph in spt_subgraphs
        ]
        feat = self.g.x_text_feat[np.concatenate(neighbors_list)]

        # nodes_pt represents the index of first node in each subgraph
        # nodes_pt[0] is the index of query node; nodes_pt[1,-1]: idx of support nodes; nodes_pt[-1]: idx of first prompt node
        nodes_pt = [0, num_nodes] + [subgraph[2] for subgraph in spt_subgraphs]
        nodes_pt = list(np.cumsum(nodes_pt))

        edge_index = torch.cat(
            [edge_index]
            + [
                spt_subgraphs[idx][1] + nodes_pt[idx + 1]
                for idx in range(len(spt_subgraphs))
            ],
            dim=1,
        )
        edge_feat = self.g.edge_text_feat.repeat([edge_index.size(1), 1])

        return feat, edge_index, edge_feat, nodes_pt

    def make_prompted_graph(self, feature_graph):
        (
            node_cls,
            feat,
            edge_index,
            edge_feat,
            nodes_pt,
            label,
            true_class,
        ) = feature_graph

        assert nodes_pt[-1] == (feat.shape[0])
        true_edge_num = len(edge_index[0])
        cls_feat = self.class_emb[node_cls]
        feat = torch.cat([feat, cls_feat])

        # Connect target node(index 0) with all class nodes
        new_node_id = 0
        qry_edge = torch.tensor(
            [
                [new_node_id] * self.n_way,
                [i + nodes_pt[-1] for i in range(self.n_way)],
            ],
            dtype=torch.long,
        )

        # Connect support nodes with corresponding class node
        spt_pt = nodes_pt[1:-1]
        cls_pt = [
            nodes_pt[-1] + i
            for i in range(self.n_way)
            for j in range(self.k_shot)
        ]
        spt_edge = torch.tensor(
            [spt_pt, cls_pt],
            dtype=torch.long,
        )

        # get final edge index and edge types
        edge_index = torch.cat(
            [
                edge_index,
                qry_edge,
                qry_edge[[1, 0]],
                spt_edge,
                spt_edge[[1, 0]],
            ],
            dim=-1,
        )
        e_type = torch.cat(
            [
                torch.zeros(true_edge_num, dtype=torch.long),
                torch.zeros(len(qry_edge[0]), dtype=torch.long) + 1,
                torch.zeros(len(qry_edge[0]), dtype=torch.long) + 2,
                torch.zeros(len(spt_edge[0]), dtype=torch.long) + 3,
                torch.zeros(len(spt_edge[0]), dtype=torch.long) + 4,
            ]
        )
        edge_feat = torch.cat(
            [
                edge_feat,
                self.g.prompt_edge_feat.repeat(
                    [(len(qry_edge[0]) + len(spt_edge[0])) * 2, 1]
                ),
            ]
        )
        assert edge_feat.size(0) == e_type.size(0)

        # get node masks
        new_subg = pyg.data.Data(
            feat, edge_index, y=label, edge_type=e_type, edge_attr=edge_feat
        )
        true_nodes_mask = torch.zeros(new_subg.num_nodes, dtype=torch.bool)
        true_nodes_mask[-self.n_way :] = True
        target_node_mask = torch.zeros(new_subg.num_nodes, dtype=torch.bool)
        target_node_mask[new_node_id] = True
        h_node_mask = torch.zeros(new_subg.num_nodes, dtype=torch.bool)
        spt_nodes_mask = torch.zeros(new_subg.num_nodes, dtype=torch.bool)
        spt_nodes_mask[nodes_pt[1:-1]] = True
        assert (spt_nodes_mask == True).sum() == self.n_way * self.k_shot

        new_subg.target_node_mask = target_node_mask
        new_subg.true_nodes_mask = true_nodes_mask
        new_subg.h_node_mask = h_node_mask
        new_subg.spt_nodes_mask = spt_nodes_mask
        new_subg.sample_num_nodes = new_subg.num_nodes
        new_subg.num_classes = self.n_way
        return new_subg

    def combine_final_graph(self, final_subgraphs):
        graph_feat = []
        graph_edge_index = []
        graph_edge_feat = []
        graph_y = []
        graph_e_type = []
        graph_target_node_mask = []
        graph_true_nodes_mask = []
        # graph_h_node_mask = []
        # graph_spt_nodes_mask = []
        pt = 0
        num_nodes = []

        for subgraph in final_subgraphs:
            graph_feat.append(subgraph.x)
            graph_edge_index.append(subgraph.edge_index + pt)
            graph_edge_feat.append(subgraph.edge_attr)
            # print(subgraph.edge_attr)
            graph_y.append(subgraph.y)
            graph_e_type.append(subgraph.edge_type)
            graph_target_node_mask.append(subgraph.target_node_mask)
            graph_true_nodes_mask.append(subgraph.true_nodes_mask)
            # graph_h_node_mask.append(subgraph.h_node_mask)
            # graph_spt_nodes_mask.append(subgraph.spt_nodes_mask)
            pt += subgraph.num_nodes
            num_nodes.append(subgraph.num_nodes)

        feat = torch.cat(graph_feat, dim=0)
        edge_index = torch.cat(graph_edge_index, dim=1)
        # print(graph_edge_feat)
        # print(len(graph_edge_feat))
        edge_feat = torch.cat(graph_edge_feat)
        y = torch.stack(graph_y, dim=0)
        e_type = torch.cat(graph_e_type, dim=0)
        target_node_mask = torch.cat(graph_target_node_mask, dim=0)
        true_nodes_mask = torch.cat(graph_true_nodes_mask, dim=0)
        # h_node_mask = torch.cat(graph_h_node_mask, dim=0)
        # spt_nodes_mask = torch.cat(graph_spt_nodes_mask, dim=0)

        new_subg = pyg.data.Data(
            feat, edge_index, y=y, edge_type=e_type, edge_attr=edge_feat
        )
        assert pt == new_subg.num_nodes
        new_subg.target_node_mask = target_node_mask
        new_subg.true_nodes_mask = true_nodes_mask
        # new_subg.h_node_mask = h_node_mask
        # new_subg.spt_nodes_mask = spt_nodes_mask
        new_subg.sample_num_nodes = new_subg.num_nodes
        # new_subg.num_classes = torch.full(
        #     (self.n_way * self.q_query,), self.n_way
        # )
        new_subg.num_classes = int(self.n_way)

        return new_subg

    def __getitem__(self, index):
        # sm = SmartTimer()
        # sm.record()
        # return node ids for an n_way k_shot q_query meta task
        # node_ids: (n_way, k_shot + q_query)
        # node_cls: (1, n_way), representing true classes corresponding to n ways
        node_ids, node_cls = next(iter(self.fs_idx_loader))

        max_n = self.n_way
        max_k = self.k_shot
        if self.kwargs["random_flag"]:
            if self.kwargs["min_n"] != self.n_way:
                self.n_way = torch.randint(
                    self.kwargs["min_n"], self.n_way, (1,)
                )[0]
            if self.kwargs["min_k"] > 0:
                if self.kwargs["min_k"] != self.k_shot:
                    self.k_shot = torch.randint(
                        self.kwargs["min_k"], self.k_shot, (1,)
                    )[0]
            else:
                self.k_shot = 0
        else:
            self.n_way = self.kwargs["min_n"]
            self.k_shot = self.kwargs["min_k"]
        node_cls = node_cls[: self.n_way]

        # spt_subgraphs will store all n_way x k_shot subgraph info
        # qry subgraphs will store all n_way x q_query subgraph info
        qry_subgraphs, spt_subgraphs, final_subgraphs = [], [], []
        for cls_idx in range(self.n_way):
            for shot_idx in range(self.k_shot + self.q_query):
                # sm.record()
                shot_subgraph = self.get_single_subgraph(
                    node_ids[cls_idx][shot_idx]
                )
                # sm.cal_and_update('get single subgraph')
                if shot_idx < self.q_query:
                    qry_subgraphs.append(shot_subgraph)
                else:
                    spt_subgraphs.append(shot_subgraph)
        assert len(qry_subgraphs) == (self.n_way * self.q_query)

        for idx, qry_subgraph in enumerate(qry_subgraphs):
            label = torch.tensor(idx // self.q_query)
            true_class = node_cls[label]
            # sm.record()
            feat, edge_index, edge_feat, nodes_pt = self.make_feature_graph(
                (qry_subgraph, spt_subgraphs)
            )
            # sm.cal_and_update('make feature graph')
            final_subgraph = self.make_prompted_graph(
                (
                    node_cls,
                    feat,
                    edge_index,
                    edge_feat,
                    nodes_pt,
                    label,
                    true_class,
                )
            )
            # sm.cal_and_update('make prompted graph')
            final_subgraphs.append(final_subgraph)

        meta_task_graph = self.combine_final_graph(final_subgraphs)
        bin_labels = torch.nn.functional.one_hot(
            meta_task_graph.y,
            self.n_way,
        ).flatten()
        total_labels = torch.zeros(meta_task_graph.num_nodes)
        total_labels[meta_task_graph.true_nodes_mask] = bin_labels.to(
            dtype=torch.float
        )
        meta_task_graph.bin_labels = total_labels
        meta_task_graph.y = torch.zeros((1, 1), dtype=torch.long)

        self.n_way = max_n
        self.k_shot = max_k

        return meta_task_graph


class FewShotHierDataset(FewShotSubgraphDataset):
    """
    FewShotHierDataset:
    true_nodes_mask: all n_way x q_query x n_way non-hierarchical prompt nodes
    target_node_mask: n_way x q_query query nodes
    spt_nodes_mask: n_way x q_query x n_way x k_shot support nodes
    h_node_mask: n_way x q_query hierarchical prompt nodes
    """

    def make_prompted_graph(self, feature_graph):
        (
            node_cls,
            feat,
            edge_index,
            edge_feat,
            nodes_pt,
            label,
            true_class,
        ) = feature_graph

        assert nodes_pt[-1] == (feat.shape[0])
        true_edge_num = len(edge_index[0])
        cls_feat = self.class_emb[node_cls]
        feat = torch.cat([feat, cls_feat, self.prompt_feat.view(1, -1)])

        new_node_id = 0

        # Connect support nodes with corresponding class node
        spt_pt = nodes_pt[1:-1]
        cls_pt = [
            nodes_pt[-1] + i
            for i in range(self.n_way)
            for j in range(self.k_shot)
        ]

        spt_edge = torch.tensor(
            [spt_pt, cls_pt],
            dtype=torch.long,
        )

        n_cls_pt = [nodes_pt[-1] + i for i in range(self.n_way)] + [
            nodes_pt[-1] + self.n_way
        ]
        h_pt = [nodes_pt[-1] + self.n_way for i in range(self.n_way)] + [
            new_node_id
        ]

        prompt_h_edge = torch.tensor(
            [
                h_pt,
                n_cls_pt,
            ],
            dtype=torch.long,
        )

        # get final edge index and edge types
        if self.single_prompt_edge:
            edge_index = torch.cat(
                [
                    edge_index,
                    spt_edge,
                    spt_edge[[1, 0]],
                    prompt_h_edge,
                ],
                dim=-1,
            )
            e_type = torch.cat(
                [
                    torch.zeros(true_edge_num, dtype=torch.long),
                    torch.zeros(len(spt_edge[0]), dtype=torch.long) + 1,
                    torch.zeros(len(spt_edge[0]), dtype=torch.long) + 2,
                    torch.zeros(len(prompt_h_edge[0]) - 1, dtype=torch.long)
                    + 3,
                    torch.zeros(1, dtype=torch.long) + 4,
                ]
            )
            edge_feat = torch.cat(
                [
                    edge_feat,
                    self.g.prompt_edge_feat.repeat(
                        [len(prompt_h_edge[0]) + len(spt_edge[0]) * 2, 1]
                    ),
                ]
            )
        else:
            edge_index = torch.cat(
                [
                    edge_index,
                    spt_edge,
                    spt_edge[[1, 0]],
                    prompt_h_edge,
                    prompt_h_edge[[1, 0]],
                ],
                dim=-1,
            )
            e_type = torch.cat(
                [
                    torch.zeros(true_edge_num, dtype=torch.long),
                    torch.zeros(len(spt_edge[0]), dtype=torch.long) + 1,
                    torch.zeros(len(spt_edge[0]), dtype=torch.long) + 2,
                    torch.zeros(len(prompt_h_edge[0]) - 1, dtype=torch.long)
                    + 3,
                    torch.zeros(1, dtype=torch.long) + 4,
                    torch.zeros(len(prompt_h_edge[0]) - 1, dtype=torch.long)
                    + 5,
                    torch.zeros(1, dtype=torch.long) + 6,
                ]
            )
            edge_feat = torch.cat(
                [
                    edge_feat,
                    self.g.prompt_edge_feat.repeat(
                        [(len(prompt_h_edge[0]) + len(spt_edge[0])) * 2, 1]
                    ),
                ]
            )
        assert edge_feat.size(0) == e_type.size(0)

        # get node masks
        new_subg = pyg.data.Data(
            feat, edge_index, y=label, edge_type=e_type, edge_attr=edge_feat
        )
        true_nodes_mask = torch.zeros(new_subg.num_nodes, dtype=torch.bool)
        true_nodes_mask[nodes_pt[-1] : -1] = True
        assert (true_nodes_mask == True).sum() == self.n_way
        h_node_mask = torch.zeros(new_subg.num_nodes, dtype=torch.bool)
        h_node_mask[-1] = True
        target_node_mask = torch.zeros(new_subg.num_nodes, dtype=torch.bool)
        target_node_mask[new_node_id] = True
        spt_nodes_mask = torch.zeros(new_subg.num_nodes, dtype=torch.bool)
        spt_nodes_mask[nodes_pt[1:-1]] = True
        assert (spt_nodes_mask == True).sum() == self.n_way * self.k_shot

        new_subg.target_node_mask = target_node_mask
        new_subg.h_node_mask = h_node_mask
        new_subg.true_nodes_mask = true_nodes_mask
        new_subg.spt_nodes_mask = spt_nodes_mask
        new_subg.sample_num_nodes = new_subg.num_nodes
        new_subg.num_classes = self.n_way
        return new_subg


class NewFewShotHierDataset(FewShotSubgraphDataset):
    def make_prompted_graph(self, feature_graph):
        (
            node_cls,
            feat,
            edge_index,
            edge_feat,
            nodes_pt,
            label,
            true_class,
        ) = feature_graph

        assert nodes_pt[-1] == (feat.shape[0])
        true_edge_num = len(edge_index[0])
        cls_feat = self.class_emb[node_cls]
        feat = torch.cat(
            [
                feat,
                cls_feat,
                self.prompt_feat.view(1, -1).repeat(1 + self.n_way, 1),
            ]
        )

        new_node_id = 0

        # Connect support nodes with corresponding class node
        spt_pt = nodes_pt[1:-1]
        cls_pt = [
            nodes_pt[-1] + i + 1 + self.n_way
            for i in range(self.n_way)
            for j in range(self.k_shot)
        ]

        spt_edge = torch.tensor(
            [spt_pt, cls_pt],
            dtype=torch.long,
        )

        # Connect spt_h_nodes with h_nodes
        spt_h_nodes = [
            nodes_pt[-1] + 1 + self.n_way + i for i in range(self.n_way)
        ]
        prompt_nodes = [nodes_pt[-1] + i for i in range(self.n_way)]
        spt_prompt_edge = torch.tensor(
            [
                spt_h_nodes,
                prompt_nodes,
            ],
            dtype=torch.long,
        )

        qry_h_edge = torch.tensor(
            [
                [nodes_pt[-1] + self.n_way for i in range(self.n_way)]
                + [new_node_id],
                [nodes_pt[-1] + i for i in range(self.n_way)]
                + [nodes_pt[-1] + self.n_way],
            ],
            dtype=torch.long,
        )

        # get final edge index and edge types
        if self.single_prompt_edge:
            edge_index = torch.cat(
                [
                    edge_index,
                    spt_edge,
                    spt_prompt_edge,
                    qry_h_edge,
                ],
                dim=-1,
            )
            e_type = torch.cat(
                [
                    torch.zeros(true_edge_num, dtype=torch.long),
                    torch.zeros(len(spt_edge[0]), dtype=torch.long) + 1,
                    torch.zeros(len(spt_prompt_edge[0]), dtype=torch.long) + 2,
                    torch.zeros(len(qry_h_edge[0]) - 1, dtype=torch.long) + 3,
                    torch.zeros(1, dtype=torch.long) + 4,
                ]
            )
            edge_feat = torch.cat(
                [
                    edge_feat,
                    self.g.prompt_edge_feat.repeat(
                        [
                            len(spt_prompt_edge[0])
                            + len(spt_edge[0])
                            + len(qry_h_edge[0]),
                            1,
                        ]
                    ),
                ]
            )
        assert edge_feat.size(0) == e_type.size(0)

        # get node masks
        new_subg = pyg.data.Data(
            feat, edge_index, y=label, edge_type=e_type, edge_attr=edge_feat
        )
        true_nodes_mask = torch.zeros(new_subg.num_nodes, dtype=torch.bool)
        true_nodes_mask[nodes_pt[-1] : nodes_pt[-1] + self.n_way] = True
        assert (true_nodes_mask == True).sum() == self.n_way
        h_node_mask = torch.zeros(new_subg.num_nodes, dtype=torch.bool)
        h_node_mask[nodes_pt[-1] + self.n_way :] = True
        target_node_mask = torch.zeros(new_subg.num_nodes, dtype=torch.bool)
        target_node_mask[new_node_id] = True
        spt_nodes_mask = torch.zeros(new_subg.num_nodes, dtype=torch.bool)
        spt_nodes_mask[nodes_pt[1:-1]] = True
        assert (spt_nodes_mask == True).sum() == self.n_way * self.k_shot

        new_subg.target_node_mask = target_node_mask
        new_subg.h_node_mask = h_node_mask
        new_subg.true_nodes_mask = true_nodes_mask
        new_subg.spt_nodes_mask = spt_nodes_mask
        new_subg.sample_num_nodes = new_subg.num_nodes
        new_subg.num_classes = self.n_way
        return new_subg


class FewShotNCDataset(FewShotSubgraphDataset):
    def make_prompted_graph(self, feature_graph):
        (
            node_cls,
            feat,
            edge_index,
            edge_feat,
            nodes_pt,
            label,
            true_class,
        ) = feature_graph

        assert nodes_pt[-1] == (feat.shape[0])
        true_edge_num = len(edge_index[0])
        cls_feat = self.class_emb[node_cls]
        feat = torch.cat(
            [
                feat,
                cls_feat,
                self.prompt_feat.view(1, -1).repeat(
                    1 + self.n_way * self.k_shot, 1
                ),
            ]
        )

        new_node_id = 0

        # Connect support nodes with corresponding class node
        spt_pt = nodes_pt[1:-1]
        spt_h_nodes = [
            nodes_pt[-1] + self.n_way + 1 + i
            for i in range(self.n_way * self.k_shot)
        ]

        spt_edge = torch.tensor(
            [
                spt_pt,
                spt_h_nodes,
            ],
            dtype=torch.long,
        )

        # Connect spt_h_nodes with h_nodes
        prompt_nodes = [
            nodes_pt[-1] + i
            for i in range(self.n_way)
            for j in range(self.k_shot)
        ]
        spt_prompt_edge = torch.tensor(
            [
                spt_h_nodes,
                prompt_nodes,
            ],
            dtype=torch.long,
        )

        # Connect qry_h_node with h_nodes, Connect qry node with qry_h_node
        qry_h_edge = torch.tensor(
            [
                [nodes_pt[-1] + self.n_way for i in range(self.n_way)]
                + [new_node_id],
                [nodes_pt[-1] + i for i in range(self.n_way)]
                + [nodes_pt[-1] + self.n_way],
            ],
            dtype=torch.long,
        )

        # get final edge index and edge types
        if self.single_prompt_edge:
            edge_index = torch.cat(
                [
                    edge_index,
                    spt_edge,
                    spt_prompt_edge,
                    qry_h_edge,
                ],
                dim=-1,
            )
            e_type = torch.cat(
                [
                    torch.zeros(true_edge_num, dtype=torch.long),
                    torch.zeros(len(spt_edge[0]), dtype=torch.long) + 1,
                    torch.zeros(len(spt_prompt_edge[0]), dtype=torch.long) + 4,
                    torch.zeros(len(qry_h_edge[0]) - 1, dtype=torch.long) + 3,
                    torch.zeros(1, dtype=torch.long) + 1,
                ]
            )
            edge_feat = torch.cat(
                [
                    edge_feat,
                    self.g.prompt_edge_feat.repeat(
                        [
                            len(spt_prompt_edge[0])
                            + len(spt_edge[0])
                            + len(qry_h_edge[0]),
                            1,
                        ]
                    ),
                ]
            )
        assert edge_feat.size(0) == e_type.size(0)

        # get node masks
        new_subg = pyg.data.Data(
            feat, edge_index, y=label, edge_type=e_type, edge_attr=edge_feat
        )
        true_nodes_mask = torch.zeros(new_subg.num_nodes, dtype=torch.bool)
        true_nodes_mask[nodes_pt[-1] : nodes_pt[-1] + self.n_way] = True
        assert (true_nodes_mask == True).sum() == self.n_way
        h_node_mask = torch.zeros(new_subg.num_nodes, dtype=torch.bool)
        h_node_mask[nodes_pt[-1] + self.n_way :] = True
        target_node_mask = torch.zeros(new_subg.num_nodes, dtype=torch.bool)
        target_node_mask[new_node_id] = True
        spt_nodes_mask = torch.zeros(new_subg.num_nodes, dtype=torch.bool)
        spt_nodes_mask[nodes_pt[1:-1]] = True
        assert (spt_nodes_mask == True).sum() == self.n_way * self.k_shot

        new_subg.target_node_mask = target_node_mask
        new_subg.h_node_mask = h_node_mask
        new_subg.true_nodes_mask = true_nodes_mask
        new_subg.spt_nodes_mask = spt_nodes_mask
        new_subg.sample_num_nodes = new_subg.num_nodes
        new_subg.num_classes = self.n_way
        return new_subg


class CLFewShotHierDataset(NewFewShotHierDataset):
    def __init__(
        self,
        pyg_graph,
        class_emb,
        data_idx,
        n_way: int,
        k_shot: int,
        q_query: int,
        drop_feature_ratio1,
        drop_feature_ratio2,
        drop_edge_ratio1,
        drop_edge_ratio2,
        datamanager: FewShotDataManager,
        mode: int,
        hop=2,
        class_mapping=None,
        prompt_feat=None,
        to_undirected=False,
        adj=None,
        single_prompt_edge=False,
        **kwargs,
    ):
        super().__init__(
            pyg_graph,
            class_emb,
            data_idx,
            n_way,
            k_shot,
            q_query,
            datamanager,
            mode,
            hop,
            class_mapping,
            prompt_feat,
            to_undirected,
            adj,
            single_prompt_edge,
            **kwargs,
        )
        self.df = [drop_feature_ratio1, drop_feature_ratio2]
        self.de = [drop_edge_ratio1, drop_edge_ratio2]

    def drop_feature(self, x: torch.Tensor, drop_prob: float) -> torch.Tensor:
        device = x.device
        drop_mask = (
            torch.empty((x.size(1),), dtype=torch.float32).uniform_(0, 1)
            < drop_prob
        )
        drop_mask = drop_mask.to(device)
        x = x.clone()
        x[:, drop_mask] = 0
        return x

    def get_aug_graph(self, node_ids, node_cls, aug_param_idx):
        qry_subgraphs, spt_subgraphs, final_subgraphs = [], [], []
        for cls_idx in range(self.n_way):
            for shot_idx in range(self.k_shot + self.q_query):
                shot_subgraph = self.get_single_subgraph(
                    node_ids[cls_idx][shot_idx]
                )
                if shot_idx < self.q_query:
                    qry_subgraphs.append(shot_subgraph)
                else:
                    spt_subgraphs.append(shot_subgraph)
        assert len(qry_subgraphs) == (self.n_way * self.q_query)
        assert len(spt_subgraphs) == (self.n_way * self.k_shot)

        for idx, qry_subgraph in enumerate(qry_subgraphs):
            label = torch.tensor(idx // self.q_query)
            true_class = node_cls[label]
            feat, edge_index, edge_feat, nodes_pt = self.make_feature_graph(
                (qry_subgraph, spt_subgraphs)
            )
            # drop edge
            edge_index, edge_mask = dropout_edge(
                edge_index, p=self.de[aug_param_idx], force_undirected=True
            )
            edge_feat = edge_feat[edge_mask]

            final_subgraph = self.make_prompted_graph(
                (
                    node_cls,
                    feat,
                    edge_index,
                    edge_feat,
                    nodes_pt,
                    label,
                    true_class,
                )
            )
            final_subgraphs.append(final_subgraph)

        return final_subgraphs

    def __getitem__(self, index):
        # return node ids for an n_way k_shot q_query meta task
        # node_ids: (n_way, k_shot + q_query)
        # node_cls: (1, n_way), representing true classes corresponding to n ways
        node_ids, node_cls = next(iter(self.fs_idx_loader))

        a1 = self.get_aug_graph(node_ids, node_cls, 0)
        a2 = self.get_aug_graph(node_ids, node_cls, 1)

        meta_task_graph1 = self.combine_final_graph(a1)
        # drop feature
        meta_task_graph1.x = self.drop_feature(meta_task_graph1.x, self.df[0])
        meta_task_graph2 = self.combine_final_graph(a2)
        # drop feature
        meta_task_graph2.x = self.drop_feature(meta_task_graph2.x, self.df[0])
        bin_labels = torch.nn.functional.one_hot(
            meta_task_graph1.y,
            self.n_way,
        ).flatten()
        meta_task_graph1.bin_labels = bin_labels.to(dtype=torch.float)
        bin_labels = torch.nn.functional.one_hot(
            meta_task_graph2.y,
            self.n_way,
        ).flatten()
        meta_task_graph2.bin_labels = bin_labels.to(dtype=torch.float)

        return meta_task_graph1, meta_task_graph2

    def get_collate_fn(self):
        def aug_collate(batch):
            # Separating the subgraphs
            subgraph1_list, subgraph2_list = list(zip(*batch))

            # Creating two separate batches
            batch1 = Batch.from_data_list(subgraph1_list)
            batch2 = Batch.from_data_list(subgraph2_list)

            return batch1, batch2

        return aug_collate


class ZeroShotNCDataset(FewShotSubgraphDataset):
    def make_feature_graph(self, graph_set):
        (qry_subgraph, _) = graph_set
        neighbors, edge_index, num_nodes = qry_subgraph
        feat = self.g.x_text_feat[neighbors.astype(int)]
        nodes_pt = [0, num_nodes]
        edge_feat = self.g.edge_text_feat.repeat([edge_index.size(1), 1])
        return feat, edge_index, edge_feat, nodes_pt

    def make_prompted_graph(self, feature_graph):
        (
            node_cls,
            feat,
            edge_index,
            edge_feat,
            nodes_pt,
            label,
            true_class,
        ) = feature_graph

        assert nodes_pt[-1] == (feat.shape[0])
        true_edge_num = len(edge_index[0])
        cls_feat = self.class_emb[node_cls]
        feat = torch.cat([feat, cls_feat, self.prompt_feat.view(1, -1)])

        new_node_id = 0

        # TODO: one direction edge. qry to hnode, hnode to prompt nodes
        h_pt = [nodes_pt[-1] + self.n_way for i in range(self.n_way)] + [
            new_node_id
        ]
        n_cls_pt = [nodes_pt[-1] + i for i in range(self.n_way)] + [
            nodes_pt[-1] + self.n_way
        ]

        prompt_h_edge = torch.tensor(
            [
                h_pt,
                n_cls_pt,
            ],
            dtype=torch.long,
        )

        # get final edge index and edge types
        edge_index = torch.cat(
            [edge_index, prompt_h_edge],
            dim=-1,
        )

        e_type = torch.cat(
            [
                torch.zeros(true_edge_num, dtype=torch.long),
                torch.zeros(len(prompt_h_edge[0]) - 1, dtype=torch.long) + 3,
                torch.zeros(1, dtype=torch.long) + 1,
            ]
        )

        edge_feat = torch.cat(
            [
                edge_feat,
                self.g.prompt_edge_feat.repeat([len(prompt_h_edge[0]), 1]),
            ]
        )

        assert edge_feat.size(0) == e_type.size(0)

        # get node masks
        new_subg = pyg.data.Data(
            feat, edge_index, y=label, edge_type=e_type, edge_attr=edge_feat
        )
        true_nodes_mask = torch.zeros(new_subg.num_nodes, dtype=torch.bool)
        true_nodes_mask[nodes_pt[-1] : -1] = True
        assert (true_nodes_mask == True).sum() == self.n_way
        h_node_mask = torch.zeros(new_subg.num_nodes, dtype=torch.bool)
        h_node_mask[-1] = True
        target_node_mask = torch.zeros(new_subg.num_nodes, dtype=torch.bool)
        target_node_mask[new_node_id] = True
        spt_nodes_mask = torch.zeros(new_subg.num_nodes, dtype=torch.bool)
        assert (spt_nodes_mask == True).sum() == self.n_way * self.k_shot

        new_subg.target_node_mask = target_node_mask
        new_subg.h_node_mask = h_node_mask
        new_subg.true_nodes_mask = true_nodes_mask
        new_subg.spt_nodes_mask = spt_nodes_mask
        new_subg.sample_num_nodes = new_subg.num_nodes
        new_subg.num_classes = self.n_way
        return new_subg


class ZeroShotSubgraphDataset(ZeroShotNCDataset):
    def make_prompted_graph(self, feature_graph):
        (
            node_cls,
            feat,
            edge_index,
            edge_feat,
            nodes_pt,
            label,
            true_class,
        ) = feature_graph

        assert nodes_pt[-1] == (feat.shape[0])
        true_edge_num = len(edge_index[0])
        cls_feat = self.class_emb[node_cls]
        feat = torch.cat([feat, cls_feat])

        new_node_id = 0

        # Add edge between all prompt nodes and query node
        n_cls_pt = [nodes_pt[-1] + i for i in range(self.n_way)]
        qry_pt = [new_node_id] * self.n_way
        prompt_edge = torch.tensor(
            [
                n_cls_pt,
                qry_pt,
            ],
            dtype=torch.long,
        )

        # get final edge index and edge types
        edge_index = torch.cat(
            [edge_index, prompt_edge, prompt_edge[[1, 0]]],
            dim=-1,
        )

        e_type = torch.cat(
            [
                torch.zeros(true_edge_num, dtype=torch.long),
                torch.zeros(len(prompt_edge[0]), dtype=torch.long) + 1,
                torch.zeros(len(prompt_edge[0]), dtype=torch.long) + 2,
            ]
        )

        edge_feat = torch.cat(
            [
                edge_feat,
                self.g.prompt_edge_feat.repeat([len(prompt_edge[0]) * 2, 1]),
            ]
        )

        assert edge_feat.size(0) == e_type.size(0)

        # get node masks
        new_subg = pyg.data.Data(
            feat, edge_index, y=label, edge_type=e_type, edge_attr=edge_feat
        )
        true_nodes_mask = torch.zeros(new_subg.num_nodes, dtype=torch.bool)
        true_nodes_mask[nodes_pt[-1] :] = True
        assert (true_nodes_mask == True).sum() == self.n_way
        h_node_mask = torch.zeros(new_subg.num_nodes, dtype=torch.bool)
        target_node_mask = torch.zeros(new_subg.num_nodes, dtype=torch.bool)
        target_node_mask[new_node_id] = True
        spt_nodes_mask = torch.zeros(new_subg.num_nodes, dtype=torch.bool)
        assert (spt_nodes_mask == True).sum() == self.n_way * self.k_shot

        new_subg.target_node_mask = target_node_mask
        new_subg.h_node_mask = h_node_mask
        new_subg.true_nodes_mask = true_nodes_mask
        new_subg.spt_nodes_mask = spt_nodes_mask
        new_subg.sample_num_nodes = new_subg.num_nodes
        new_subg.num_classes = self.n_way
        return new_subg


class CLZeroShotNCDataset(ZeroShotNCDataset):
    def __init__(
        self,
        pyg_graph,
        class_emb,
        data_idx,
        n_way: int,
        k_shot: int,
        q_query: int,
        drop_feature_ratio1,
        drop_feature_ratio2,
        drop_edge_ratio1,
        drop_edge_ratio2,
        datamanager: FewShotDataManager,
        mode: int,
        hop=2,
        class_mapping=None,
        prompt_feat=None,
        to_undirected=False,
        adj=None,
        **kwargs,
    ):
        super().__init__(
            pyg_graph,
            class_emb,
            data_idx,
            n_way,
            k_shot,
            q_query,
            datamanager,
            mode,
            hop,
            class_mapping,
            prompt_feat,
            to_undirected,
            adj,
            **kwargs,
        )
        self.df = [drop_feature_ratio1, drop_feature_ratio2]
        self.de = [drop_edge_ratio1, drop_edge_ratio2]

    def drop_feature(self, x: torch.Tensor, drop_prob: float) -> torch.Tensor:
        device = x.device
        drop_mask = (
            torch.empty((x.size(1),), dtype=torch.float32).uniform_(0, 1)
            < drop_prob
        )
        drop_mask = drop_mask.to(device)
        x = x.clone()
        x[:, drop_mask] = 0
        return x

    def get_aug_graph(self, node_ids, node_cls, aug_param_idx):
        qry_subgraphs, spt_subgraphs, final_subgraphs = [], [], []
        for cls_idx in range(self.n_way):
            for shot_idx in range(self.k_shot + self.q_query):
                shot_subgraph = self.get_single_subgraph(
                    node_ids[cls_idx][shot_idx]
                )
                if shot_idx < self.q_query:
                    qry_subgraphs.append(shot_subgraph)
                else:
                    spt_subgraphs.append(shot_subgraph)
        assert len(qry_subgraphs) == (self.n_way * self.q_query)
        assert len(spt_subgraphs) == (self.n_way * self.k_shot)

        for idx, qry_subgraph in enumerate(qry_subgraphs):
            label = torch.tensor(idx // self.q_query)
            true_class = node_cls[label]
            feat, edge_index, edge_feat, nodes_pt = self.make_feature_graph(
                (qry_subgraph, spt_subgraphs)
            )
            # drop edge
            edge_index, edge_mask = dropout_edge(
                edge_index, p=self.de[aug_param_idx], force_undirected=True
            )
            edge_feat = edge_feat[edge_mask]

            final_subgraph = self.make_prompted_graph(
                (
                    node_cls,
                    feat,
                    edge_index,
                    edge_feat,
                    nodes_pt,
                    label,
                    true_class,
                )
            )
            final_subgraphs.append(final_subgraph)

        return final_subgraphs

    def __getitem__(self, index):
        # return node ids for an n_way k_shot q_query meta task
        # node_ids: (n_way, k_shot + q_query)
        # node_cls: (1, n_way), representing true classes corresponding to n ways
        node_ids, node_cls = next(iter(self.fs_idx_loader))

        a1 = self.get_aug_graph(node_ids, node_cls, 0)
        a2 = self.get_aug_graph(node_ids, node_cls, 1)

        meta_task_graph1 = self.combine_final_graph(a1)
        # drop feature
        meta_task_graph1.x = self.drop_feature(meta_task_graph1.x, self.df[0])
        meta_task_graph2 = self.combine_final_graph(a2)
        # drop feature
        meta_task_graph2.x = self.drop_feature(meta_task_graph2.x, self.df[0])
        bin_labels = torch.nn.functional.one_hot(
            meta_task_graph1.y,
            self.n_way,
        ).flatten()
        meta_task_graph1.bin_labels = bin_labels.to(dtype=torch.float)
        bin_labels = torch.nn.functional.one_hot(
            meta_task_graph2.y,
            self.n_way,
        ).flatten()
        meta_task_graph2.bin_labels = bin_labels.to(dtype=torch.float)

        return meta_task_graph1, meta_task_graph2

    def get_collate_fn(self):
        def aug_collate(batch):
            # Separating the subgraphs
            subgraph1_list, subgraph2_list = list(zip(*batch))

            # Creating two separate batches
            batch1 = Batch.from_data_list(subgraph1_list)
            batch2 = Batch.from_data_list(subgraph2_list)

            return batch1, batch2

        return aug_collate


class FewShotKGHierDataset(FewShotSubgraphDataset):
    def __init__(
        self,
        pyg_graph,
        class_emb,
        data_idx,
        n_way,
        k_shot,
        q_query,
        datamanager: FewShotDataManager,
        mode,
        edges,  # all edges in the graph
        fs_edges,  # edges that belongs to specific classes for few-shot
        fs_edge_types,
        hop=2,
        prompt_feat=None,
        to_undirected=False,
        single_prompt_edge=True,
        adj=None,
        **kwargs,
    ):
        self.g = pyg_graph
        self.kwargs = kwargs
        # mode 0 for sample index from training classes, 1 for val, 2 for test
        self.fs_idx_loader = datamanager.get_data_loader(mode)
        self.n_way = n_way
        self.k_shot = k_shot
        self.q_query = q_query
        self.edges = edges
        self.fs_edges = fs_edges
        self.fs_edge_types = fs_edge_types
        self.prompt_feat = prompt_feat
        self.single_prompt_edge = single_prompt_edge

        edge_index = fs_edges
        if to_undirected:
            edge_index = pyg.utils.to_undirected(edge_index)

        if adj is not None:
            self.adj = adj
        else:
            self.adj = csr_array(
                (
                    torch.ones(len(edge_index[0])),
                    (edge_index[0], edge_index[1]),
                ),
                shape=(self.g.num_nodes, self.g.num_nodes),
            )
        self.class_emb = class_emb
        self.hop = hop
        self.data_idx = data_idx

    def get_single_subgraph(self, node_id):
        node_ids = self.edges[:, node_id]
        neighbors = sample_fixed_hop_size_neighbor(
            self.adj, node_ids, self.hop, max_nodes_per_hope=100
        )
        neighbors = np.r_[node_ids, neighbors]

        node_mask = torch.zeros(self.g.num_nodes, dtype=torch.bool)
        node_mask[neighbors] = True

        edge_mask = node_mask[self.fs_edges[0]] & node_mask[self.fs_edges[1]]

        edge2idx = torch.zeros(self.g.num_nodes, dtype=torch.long)
        edge2idx[neighbors] = torch.arange(len(neighbors))
        edge_index = self.fs_edges[:, edge_mask]
        edge_type = self.fs_edge_types[edge_mask]
        edge_index = edge2idx[edge_index]

        return neighbors, edge_index, neighbors.shape[0], edge_type

    def make_feature_graph(self, graph_set):
        """
        For each query node, combine its subgraph with all support subgraphs.
        Args:
            qry_subgraph:
            spt_subgraphs:

        Returns:

        """
        (qry_subgraph, spt_subgraphs) = graph_set
        neighbors, edge_index, num_nodes, edge_type = qry_subgraph
        neighbors_list = [neighbors] + [
            subgraph[0] for subgraph in spt_subgraphs
        ]
        feat = self.g.x_text_feat[np.concatenate(neighbors_list).astype(int)]
        if not isinstance(feat, torch.Tensor):
            feat = torch.from_numpy(feat)
            feat = feat.float()

        # nodes_pt represents the index of first node in each subgraph
        # nodes_pt[0] is the index of query node; nodes_pt[1,-1]: idx of support nodes; nodes_pt[-1]: idx of first prompt node
        nodes_pt = [0, num_nodes] + [subgraph[2] for subgraph in spt_subgraphs]
        nodes_pt = list(np.cumsum(nodes_pt))

        edge_index = torch.cat(
            [edge_index]
            + [
                spt_subgraphs[idx][1] + nodes_pt[idx + 1]
                for idx in range(len(spt_subgraphs))
            ],
            dim=1,
        )
        edge_index = torch.cat([edge_index, edge_index[[1, 0]]], dim=-1)
        edge_type = torch.cat(
            [edge_type]
            + [spt_subgraphs[idx][3] for idx in range(len(spt_subgraphs))],
        )

        edge_feat = self.g.edge_text_feat[
            torch.cat(
                [edge_type, edge_type + int(len(self.g.edge_text_feat) / 2)]
            )
        ]

        return feat, edge_index, edge_feat, nodes_pt

    def make_prompted_graph(self, feature_graph):
        (
            node_cls,
            feat,
            edge_index,
            edge_feat,
            nodes_pt,
            label,
            true_class,
        ) = feature_graph

        assert nodes_pt[-1] == (feat.shape[0])
        true_edge_num = len(edge_index[0])
        cls_feat = self.class_emb[node_cls]
        feat = torch.cat(
            [
                feat,
                cls_feat,
                self.prompt_feat.view(1, -1).repeat(1 + self.n_way, 1),
            ]
        )

        new_node_id = 0

        # Connect support nodes with corresponding class node
        spt_pt = nodes_pt[1:-1]
        cls_pt = [
            nodes_pt[-1] + i + 1 + self.n_way
            for i in range(self.n_way)
            for j in range(self.k_shot)
        ]

        spt_edge = torch.tensor(
            [
                spt_pt + [pt + 1 for pt in spt_pt],
                cls_pt * 2,
            ],
            dtype=torch.long,
        )

        # Connect spt_h_nodes with h_nodes
        spt_h_nodes = [
            nodes_pt[-1] + 1 + self.n_way + i for i in range(self.n_way)
        ]
        prompt_nodes = [nodes_pt[-1] + i for i in range(self.n_way)]
        spt_prompt_edge = torch.tensor(
            [
                spt_h_nodes,
                prompt_nodes,
            ],
            dtype=torch.long,
        )

        # Connect qry_h_node
        qry_h_edge = torch.tensor(
            [
                [nodes_pt[-1] + self.n_way for i in range(self.n_way)]
                + [new_node_id, new_node_id + 1],
                [nodes_pt[-1] + i for i in range(self.n_way)]
                + [nodes_pt[-1] + self.n_way] * 2,
            ],
            dtype=torch.long,
        )

        # get final edge index and edge types
        if self.single_prompt_edge:
            edge_index = torch.cat(
                [
                    edge_index,
                    spt_edge,
                    spt_prompt_edge,
                    qry_h_edge,
                ],
                dim=-1,
            )
            e_type = torch.cat(
                [
                    torch.zeros(true_edge_num, dtype=torch.long),
                    torch.zeros(len(spt_edge[0]), dtype=torch.long) + 1,
                    torch.zeros(len(spt_prompt_edge[0]), dtype=torch.long) + 2,
                    torch.zeros(len(qry_h_edge[0]) - 2, dtype=torch.long) + 3,
                    torch.zeros(2, dtype=torch.long) + 4,
                ]
            )
            edge_feat = torch.cat(
                [
                    edge_feat,
                    self.g.prompt_edge_feat.repeat(
                        [
                            len(spt_prompt_edge[0])
                            + len(spt_edge[0])
                            + len(qry_h_edge[0]),
                            1,
                        ]
                    ),
                ]
            )
        assert edge_feat.size(0) == e_type.size(0)

        # get node masks
        new_subg = pyg.data.Data(
            feat, edge_index, y=label, edge_type=e_type, edge_attr=edge_feat
        )
        true_nodes_mask = torch.zeros(new_subg.num_nodes, dtype=torch.bool)
        true_nodes_mask[nodes_pt[-1] : nodes_pt[-1] + self.n_way] = True
        assert (true_nodes_mask == True).sum() == self.n_way
        h_node_mask = torch.zeros(new_subg.num_nodes, dtype=torch.bool)
        h_node_mask[nodes_pt[-1] + self.n_way :] = True
        target_node_mask = torch.zeros(new_subg.num_nodes, dtype=torch.bool)
        target_node_mask[new_node_id] = True
        spt_nodes_mask = torch.zeros(new_subg.num_nodes, dtype=torch.bool)
        spt_nodes_mask[nodes_pt[1:-1]] = True
        assert (spt_nodes_mask == True).sum() == self.n_way * self.k_shot

        new_subg.target_node_mask = target_node_mask
        new_subg.h_node_mask = h_node_mask
        new_subg.true_nodes_mask = true_nodes_mask
        new_subg.spt_nodes_mask = spt_nodes_mask
        new_subg.sample_num_nodes = new_subg.num_nodes
        new_subg.num_classes = self.n_way
        return new_subg


class FewShotKGDataset(FewShotKGHierDataset):
    def make_prompted_graph(self, feature_graph):
        (
            node_cls,
            feat,
            edge_index,
            edge_feat,
            nodes_pt,
            label,
            true_class,
        ) = feature_graph

        assert nodes_pt[-1] == (feat.shape[0])
        true_edge_num = len(edge_index[0])
        cls_feat = self.class_emb[node_cls]
        feat = torch.cat(
            [
                feat,
                cls_feat,
                self.prompt_feat.view(1, -1).repeat(
                    1 + self.n_way * self.k_shot, 1
                ),
            ]
        )

        new_node_id = 0

        # Connect support nodes with corresponding class node
        spt_pt = nodes_pt[1:-1]

        cls_pt = [
            nodes_pt[-1] + i + 1 + self.n_way
            for i in range(self.n_way)
            for j in range(self.k_shot)
        ]

        spt_edge = torch.tensor(
            [
                spt_pt + [pt + 1 for pt in spt_pt],
                cls_pt * 2,
            ],
            dtype=torch.long,
        )

        # Connect spt_h_nodes with h_nodes
        spt_h_nodes = [
            nodes_pt[-1] + 1 + self.n_way + i for i in range(self.n_way)
        ]
        prompt_nodes = [nodes_pt[-1] + i for i in range(self.n_way)]
        spt_prompt_edge = torch.tensor(
            [
                spt_h_nodes,
                prompt_nodes,
            ],
            dtype=torch.long,
        )

        # Connect qry_h_node
        qry_h_edge = torch.tensor(
            [
                [nodes_pt[-1] + self.n_way for i in range(self.n_way)]
                + [new_node_id, new_node_id + 1]
                + [nodes_pt[-1] + self.n_way] * 2,
                [nodes_pt[-1] + i for i in range(self.n_way)]
                + [nodes_pt[-1] + self.n_way] * 2
                + [new_node_id, new_node_id + 1],
            ],
            dtype=torch.long,
        )

        # get final edge index and edge types
        if self.single_prompt_edge:
            edge_index = torch.cat(
                [
                    edge_index,
                    spt_edge,
                    spt_edge[[1, 0]],
                    spt_prompt_edge,
                    qry_h_edge,
                ],
                dim=-1,
            )
            e_type = torch.cat(
                [
                    torch.zeros(true_edge_num, dtype=torch.long),
                    torch.zeros(len(spt_edge[0]), dtype=torch.long) + 1,
                    torch.zeros(len(spt_edge[0]), dtype=torch.long) + 2,
                    torch.zeros(len(spt_prompt_edge[0]), dtype=torch.long) + 4,
                    torch.zeros(len(qry_h_edge[0]) - 4, dtype=torch.long) + 3,
                    torch.zeros(2, dtype=torch.long) + 1,
                    torch.zeros(2, dtype=torch.long) + 2,
                ]
            )
            edge_feat = torch.cat(
                [
                    edge_feat,
                    self.g.prompt_edge_feat.repeat(
                        [
                            len(spt_prompt_edge[0])
                            + len(spt_edge[0]) * 2
                            + len(qry_h_edge[0]),
                            1,
                        ]
                    ),
                ]
            )
        assert edge_feat.size(0) == e_type.size(0)

        # get node masks
        new_subg = pyg.data.Data(
            feat, edge_index, y=label, edge_type=e_type, edge_attr=edge_feat
        )
        true_nodes_mask = torch.zeros(new_subg.num_nodes, dtype=torch.bool)
        true_nodes_mask[nodes_pt[-1] : nodes_pt[-1] + self.n_way] = True
        assert (true_nodes_mask == True).sum() == self.n_way
        h_node_mask = torch.zeros(new_subg.num_nodes, dtype=torch.bool)
        h_node_mask[nodes_pt[-1] + self.n_way :] = True
        target_node_mask = torch.zeros(new_subg.num_nodes, dtype=torch.bool)
        target_node_mask[new_node_id] = True
        spt_nodes_mask = torch.zeros(new_subg.num_nodes, dtype=torch.bool)
        spt_nodes_mask[nodes_pt[1:-1]] = True
        assert (spt_nodes_mask == True).sum() == self.n_way * self.k_shot

        new_subg.target_node_mask = target_node_mask
        new_subg.h_node_mask = h_node_mask
        new_subg.true_nodes_mask = true_nodes_mask
        new_subg.spt_nodes_mask = spt_nodes_mask
        new_subg.sample_num_nodes = new_subg.num_nodes
        new_subg.num_classes = self.n_way
        return new_subg


class ZeroShotKGDataset(FewShotKGHierDataset):
    def make_feature_graph(self, graph_set):
        (qry_subgraph, _) = graph_set
        neighbors, edge_index, num_nodes, edge_type = qry_subgraph
        feat = self.g.x_text_feat[neighbors.astype(int)]
        if not isinstance(feat, torch.Tensor):
            feat = torch.from_numpy(feat)
            feat = feat.float()

        nodes_pt = [0, num_nodes]
        edge_feat = self.g.edge_text_feat[
            torch.cat(
                [edge_type, edge_type + int(len(self.g.edge_text_feat) / 2)]
            )
        ]
        edge_index = torch.cat([edge_index, edge_index[[1, 0]]], dim=-1)

        return feat, edge_index, edge_feat, nodes_pt

    def make_prompted_graph(self, feature_graph):
        (
            node_cls,
            feat,
            edge_index,
            edge_feat,
            nodes_pt,
            label,
            true_class,
        ) = feature_graph

        assert nodes_pt[-1] == (feat.shape[0])
        true_edge_num = len(edge_index[0])
        cls_feat = self.class_emb[node_cls]
        feat = torch.cat([feat, cls_feat, self.prompt_feat.view(1, -1)])

        new_node_id = 0

        # Connect qry_h_node
        qry_h_edge = torch.tensor(
            [
                [nodes_pt[-1] + self.n_way for i in range(self.n_way)]
                + [new_node_id, new_node_id + 1]
                + [nodes_pt[-1] + self.n_way] * 2,
                [nodes_pt[-1] + i for i in range(self.n_way)]
                + [nodes_pt[-1] + self.n_way] * 2
                + [new_node_id, new_node_id + 1],
            ],
            dtype=torch.long,
        )

        # get final edge index and edge types
        if self.single_prompt_edge:
            edge_index = torch.cat(
                [
                    edge_index,
                    qry_h_edge,
                ],
                dim=-1,
            )
            e_type = torch.cat(
                [
                    torch.zeros(true_edge_num, dtype=torch.long),
                    torch.zeros(len(qry_h_edge[0]) - 4, dtype=torch.long) + 3,
                    torch.zeros(2, dtype=torch.long) + 1,
                    torch.zeros(2, dtype=torch.long) + 2,
                ]
            )
            edge_feat = torch.cat(
                [
                    edge_feat,
                    self.g.prompt_edge_feat.repeat([len(qry_h_edge[0]), 1]),
                ]
            )
        assert edge_feat.size(0) == e_type.size(0)

        # get node masks
        new_subg = pyg.data.Data(
            feat, edge_index, y=label, edge_type=e_type, edge_attr=edge_feat
        )
        true_nodes_mask = torch.zeros(new_subg.num_nodes, dtype=torch.bool)
        true_nodes_mask[nodes_pt[-1] : nodes_pt[-1] + self.n_way] = True
        assert (true_nodes_mask == True).sum() == self.n_way
        h_node_mask = torch.zeros(new_subg.num_nodes, dtype=torch.bool)
        h_node_mask[nodes_pt[-1] + self.n_way :] = True
        target_node_mask = torch.zeros(new_subg.num_nodes, dtype=torch.bool)
        target_node_mask[new_node_id] = True
        spt_nodes_mask = torch.zeros(new_subg.num_nodes, dtype=torch.bool)
        assert (spt_nodes_mask == True).sum() == self.n_way * self.k_shot

        new_subg.target_node_mask = target_node_mask
        new_subg.h_node_mask = h_node_mask
        new_subg.true_nodes_mask = true_nodes_mask
        new_subg.spt_nodes_mask = spt_nodes_mask
        new_subg.sample_num_nodes = new_subg.num_nodes
        new_subg.num_classes = self.n_way
        return new_subg


class MultiDataset(DatasetWithCollate):
    def __init__(
        self,
        datas,
        dataset_multiple=1,
        window_size=3,
        patience=3,
        min_ratio=0.1,
        # walk_length=None,
    ):
        self.datas = datas
        self.sizes = np.array([len(d) for d in datas])
        self.performance_record = []
        self.patience = patience
        if isinstance(self.patience, int):
            self.patience = np.zeros(len(self.sizes)) + self.patience
        self.inpatience = np.zeros(len(self.patience))
        self.window_size = window_size
        if isinstance(self.window_size, int):
            self.window_size = np.zeros(len(self.sizes)) + self.window_size
        self.dataset_multiple = dataset_multiple
        if not isinstance(self.dataset_multiple, list):
            self.dataset_multiple = (
                np.zeros(len(self.sizes)) + self.dataset_multiple
            )
        self.min_ratio = min_ratio
        if isinstance(self.min_ratio, float):
            self.min_ratio = np.zeros(len(self.sizes)) + self.min_ratio
        # self.walk_length = walk_length
        self.compute_sizes()

    def compute_sizes(self):
        self.aug_sizes = (self.sizes * np.array(self.dataset_multiple)).astype(
            int
        )
        self.size_seg = np.cumsum(self.aug_sizes)
        self.ind2dataset = np.arange(len(self.datas)).repeat(self.aug_sizes)
        self.sample_ind = (
            np.random.rand(len(self.ind2dataset))
            * self.sizes.repeat(self.aug_sizes)
        ).astype(int)
        self.data_start_index = np.r_[0, self.size_seg[:-1]]

    def __len__(self):
        return np.sum(self.aug_sizes)

    def __getitem__(self, index):
        dataset_ind = self.ind2dataset[index]
        dataset = self.datas[dataset_ind]
        ret_data = dataset[self.sample_ind[index]]
        # if self.walk_length is not None:
        #     ret_data.rwpe = scipy_rwpe(ret_data, self.walk_length)
        #     print(ret_data.rwpe)
        return ret_data

    def get_collate_fn(self):
        return self.datas[0].get_collate_fn()

    def update(self, metric):
        for i in range(len(self.datas)):
            if len(self.performance_record) < self.window_size[i]:
                continue
            vals = [entry[i] for entry in self.performance_record][
                -int(self.window_size[i]) :
            ]
            vals = np.array(vals)
            mean = vals.mean()
            if metric[i] > mean:
                self.inpatience[i] = 0
            else:
                self.inpatience[i] += 1
            if self.inpatience[i] > self.patience[i]:
                self.dataset_multiple[i] = max(
                    self.min_ratio[i], self.dataset_multiple[i] / 2
                )
                # self.inpatience[i] = 0
        self.compute_sizes()
        self.performance_record.append(metric)

ofa_datasets.py
LechengKong/OneForAll
import argparse
import os
import torch
import torch_geometric as pyg
from pytorch_lightning.loggers import WandbLogger
from ofa_datasets import (
    GraphListDataset,
    SubgraphDataset,
    MultiDataset,
    GraphListHierDataset,
    SubgraphHierDataset,
    GraphListHierFSDataset,
    GraphListHierFixDataset,
    SubgraphLinkHierDataset,
    SubgraphKGHierDataset,
    FewShotNCDataset,
    FewShotKGDataset,
    ZeroShotNCDataset,
    ZeroShotKGDataset,
    SubgraphNopromptDataset,
    GraphListNopromptDataset,
)
from fs_datamanager import FewShotDataManager
from gp.utils.utils import (
    load_yaml,
    combine_dict,
    merge_mod,
    setup_exp,
    set_random_seed,
    k_fold2_split,
    k_fold_ind,
)
from gp.lightning.metric import (
    binary_auc_func,
    flat_binary_func,
    classification_func,
    EvalKit,
)
from gp.lightning.data_template import DataModule, DataWithMeta
from gp.lightning.training import lightning_fit, lightning_test
from gp.lightning.module_template import ExpConfig
from types import SimpleNamespace
from lightning_model import GraphPredLightning
from models.model import BinGraphModel, BinGraphAttModel, AdaPoolClassModel
from gp.nn.models.pyg import PyGGIN, PyGRGCN, PyGGINE
from models.model import PyGRGCNEdge

from torchmetrics import AUROC, Accuracy
from utils import (
    SentenceEncoder,
    binary_apr_func,
    MultiApr,
    flat_label_func,
    label_apr_func,
    binary_single_auc_func,
    binary_auc_multi_func,
    MultiAuc,
)
from data.arxiv.gen_data import ArxivOFADataset
from data.Cora.gen_data import CoraOFADataset
from data.Pubmed.gen_data import PubmedOFADataset
from data.WN18RR.gen_data import WN18RROFADataset
from data.FB15K237.gen_data import FB15K237OFADataset
from data.wikics.gen_data import WikiCSOFADataset
from data.chemblpre.gen_data import CHEMBLPREOFADataset
from data.chempcba.gen_data import CHEMPCBAOFADataset
from data.chemhiv.gen_data import CHEMHIVOFADataset

from gp.lightning.metric import flat_binary_func_fs
from gp.utils.utils import SmartTimer
from scipy.sparse import csr_array


def data_construct(
    data_names,
    encoder,
    batch_size=None,
    sample_size=None,
    walk_length=None,
    fs_setting=None,
):
    constructed_data = []
    for data_idx, name in enumerate(data_names):
        print("loading: ", name)
        if name == "arxiv":
            dataset = ArxivOFADataset("arxiv", sentence_encoder=encoder)
            text_g = dataset.data
            text_g.x = text_g.x_text_feat
            text_g.prompt_edge_feat = dataset.prompt_edge_feat
            kfold = k_fold_ind(text_g.y, 10)
            text_split = k_fold2_split(kfold, len(text_g.y))[0]

            def trim_class(label, num_class):
                binary_rep = torch.zeros((1, num_class))
                binary_rep[0, label] = 1
                return label.view(1, -1), binary_rep

            def make_data(split_name, state_name):
                return DataWithMeta(
                    SubgraphHierDataset(
                        text_g,
                        text_g.label_text_feat,
                        text_split[split_name],
                        prompt_feat=text_g.prompt_text_feat,
                        to_undirected=True,
                        trim_class_func=trim_class,
                        walk_length=walk_length,
                    ),
                    batch_size,
                    sample_size=sample_size,
                    metric="acc",
                    state_name=state_name,
                    classes=40,
                    meta_data={"eval_func": classification_func},
                )

            split_data = {
                "train": [
                    SubgraphHierDataset(
                        text_g,
                        text_g.label_text_feat,
                        text_split[0],
                        prompt_feat=text_g.prompt_text_feat,
                        to_undirected=True,
                        trim_class_func=trim_class,
                        walk_length=walk_length,
                    )
                ],
                "test": [
                    make_data(2, "test_arxiv"),
                    make_data(0, "test_train_arxiv"),
                ],
                "val": [make_data(1, "valid_arxiv")],
            }
            constructed_data.append(split_data)

        if name == "arxivnohier":
            dataset = ArxivOFADataset("arxiv", sentence_encoder=encoder)
            text_g = dataset.data
            text_g.x = text_g.x_text_feat
            text_g.prompt_edge_feat = dataset.prompt_edge_feat
            kfold = k_fold_ind(text_g.y, 10)
            text_split = k_fold2_split(kfold, len(text_g.y))[0]

            def trim_class(label, num_class):
                binary_rep = torch.zeros((1, num_class))
                binary_rep[0, label] = 1
                return label.view(1, -1), binary_rep

            def make_data(split_name, state_name):
                return DataWithMeta(
                    SubgraphDataset(
                        text_g,
                        text_g.label_text_feat,
                        text_split[split_name],
                        to_undirected=True,
                        trim_class_func=trim_class,
                        walk_length=walk_length,
                    ),
                    batch_size,
                    sample_size=sample_size,
                    metric="acc",
                    state_name=state_name,
                    classes=40,
                    meta_data={"eval_func": classification_func},
                )

            split_data = {
                "train": [
                    SubgraphDataset(
                        text_g,
                        text_g.label_text_feat,
                        text_split[0],
                        to_undirected=True,
                        trim_class_func=trim_class,
                        walk_length=walk_length,
                    )
                ],
                "test": [
                    make_data(2, "test_arxiv"),
                    make_data(0, "test_train_arxiv"),
                ],
                "val": [make_data(1, "valid_arxiv")],
            }
            constructed_data.append(split_data)
        if name == "arxivnoprompt":
            dataset = ArxivOFADataset("arxiv", sentence_encoder=encoder)
            text_g = dataset.data
            text_g.x = text_g.x_text_feat
            text_g.prompt_edge_feat = dataset.prompt_edge_feat
            kfold = k_fold_ind(text_g.y, 10)
            text_split = k_fold2_split(kfold, len(text_g.y))[0]

            def trim_class(label, num_class):
                binary_rep = torch.zeros((1, num_class))
                binary_rep[0, label] = 1
                return label.view(1, -1), binary_rep

            def make_data(split_name, state_name):
                return DataWithMeta(
                    SubgraphNopromptDataset(
                        text_g,
                        text_g.label_text_feat,
                        text_split[split_name],
                        to_undirected=True,
                        trim_class_func=trim_class,
                        walk_length=walk_length,
                    ),
                    batch_size,
                    sample_size=sample_size,
                    metric="acc",
                    state_name=state_name,
                    classes=40,
                    meta_data={"eval_func": classification_func},
                )

            split_data = {
                "train": [
                    SubgraphNopromptDataset(
                        text_g,
                        text_g.label_text_feat,
                        text_split[0],
                        to_undirected=True,
                        trim_class_func=trim_class,
                        walk_length=walk_length,
                    )
                ],
                "test": [
                    make_data(2, "test_arxiv"),
                    make_data(0, "test_train_arxiv"),
                ],
                "val": [make_data(1, "valid_arxiv")],
            }
            constructed_data.append(split_data)

        if name in ["chempcba", "chemblpre"]:
            if name == "chempcba":
                dataset = CHEMPCBAOFADataset(
                    "chempcba", sentence_encoder=encoder
                )
            elif name == "chemblpre":
                dataset = CHEMBLPREOFADataset(
                    "chemblpre", sentence_encoder=encoder
                )
            # print(dataset[350342])
            # return 0
            split = dataset.get_idx_split()

            def trim_class(embs, classes):
                valid_idx = classes == classes
                # valid_idx = torch.zeros_like(classes, dtype=torch.bool)
                return (
                    torch.tensor([[0]]),
                    embs[valid_idx.view(-1)].detach().clone(),
                    classes[:, valid_idx.view(-1)].detach().clone(),
                )

            def make_data(split_name, state_name):
                return DataWithMeta(
                    GraphListHierDataset(
                        dataset,
                        dataset.label_text_feat,
                        dataset.prompt_edge_feat,
                        dataset.prompt_text_feat,
                        split[split_name],
                        single_prompt_edge=True,
                        # trim_class_func=trim_class,
                        walk_length=walk_length,
                    ),
                    batch_size,
                    sample_size=sample_size,
                    metric="apr",
                    state_name=state_name,
                    classes=len(dataset.label_text_feat),
                    meta_data={"eval_func": binary_apr_func},
                )

            split_data = {
                "train": [
                    GraphListHierDataset(
                        dataset,
                        dataset.label_text_feat,
                        dataset.prompt_edge_feat,
                        dataset.prompt_text_feat,
                        split["train"],
                        trim_class_func=trim_class,
                        single_prompt_edge=True,
                        walk_length=walk_length,
                    )
                ],
                "test": [
                    make_data("test", "test_" + name),
                    make_data("train", "test_train_" + name),
                ],
                "val": [make_data("valid", "valid_" + name)],
            }
            constructed_data.append(split_data)
        if name in ["chemhiv"]:
            if name == "chemhiv":
                dataset = CHEMHIVOFADataset(
                    "chemhiv", sentence_encoder=encoder
                )
            # print(dataset[350342])
            # return 0
            split = dataset.get_idx_split()

            # def trim_class(embs, classes):
            #     valid_idx = classes == classes
            #     # valid_idx = torch.zeros_like(classes, dtype=torch.bool)
            #     return (
            #         torch.tensor([[0]]),
            #         embs[valid_idx.view(-1)].detach().clone(),
            #         classes[:, valid_idx.view(-1)].detach().clone(),
            #     )

            def trim_class(embs, label):
                label = label.to(torch.long)
                one_hot_label = torch.nn.functional.one_hot(
                    label, num_classes=2
                )
                return label, embs, one_hot_label

            # print(sub_dataset.label_text_feat.size())

            def make_data(split_name, state_name):
                return DataWithMeta(
                    GraphListHierDataset(
                        dataset,
                        dataset.label_text_feat,
                        dataset.prompt_edge_feat,
                        dataset.prompt_text_feat,
                        split[split_name],
                        single_prompt_edge=True,
                        # trim_class_func=trim_class,
                        walk_length=walk_length,
                    ),
                    batch_size,
                    sample_size=sample_size,
                    metric="auc",
                    state_name=state_name,
                    classes=len(dataset.label_text_feat),
                    meta_data={"eval_func": binary_auc_func},
                )

            split_data = {
                "train": [
                    GraphListHierDataset(
                        dataset,
                        dataset.label_text_feat,
                        dataset.prompt_edge_feat,
                        dataset.prompt_text_feat,
                        split["train"],
                        trim_class_func=trim_class,
                        single_prompt_edge=True,
                        walk_length=walk_length,
                    )
                ],
                "test": [
                    make_data("test", "test_" + name),
                    make_data("train", "test_train_" + name),
                ],
                "val": [make_data("valid", "valid_" + name)],
            }
            constructed_data.append(split_data)

        if name in ["chemhivnohier"]:
            if name == "chemhivnohier":
                dataset = CHEMHIVOFADataset(
                    "chemhiv", sentence_encoder=encoder
                )
            # print(dataset[350342])
            # return 0
            split = dataset.get_idx_split()

            # def trim_class(embs, classes):
            #     valid_idx = classes == classes
            #     # valid_idx = torch.zeros_like(classes, dtype=torch.bool)
            #     return (
            #         torch.tensor([[0]]),
            #         embs[valid_idx.view(-1)].detach().clone(),
            #         classes[:, valid_idx.view(-1)].detach().clone(),
            #     )

            def trim_class(embs, label):
                label = label.to(torch.long)
                one_hot_label = torch.nn.functional.one_hot(
                    label, num_classes=2
                )
                return label, embs, one_hot_label

            # print(sub_dataset.label_text_feat.size())

            def make_data(split_name, state_name):
                return DataWithMeta(
                    GraphListDataset(
                        dataset,
                        dataset.label_text_feat,
                        dataset.prompt_edge_feat,
                        split[split_name],
                        single_prompt_edge=True,
                        # trim_class_func=trim_class,
                        walk_length=walk_length,
                    ),
                    batch_size,
                    sample_size=sample_size,
                    metric="auc",
                    state_name=state_name,
                    classes=len(dataset.label_text_feat),
                    meta_data={"eval_func": binary_auc_func},
                )

            split_data = {
                "train": [
                    GraphListDataset(
                        dataset,
                        dataset.label_text_feat,
                        dataset.prompt_edge_feat,
                        split["train"],
                        trim_class_func=trim_class,
                        single_prompt_edge=True,
                        walk_length=walk_length,
                    )
                ],
                "test": [
                    make_data("test", "test_" + name),
                    make_data("train", "test_train_" + name),
                ],
                "val": [make_data("valid", "valid_" + name)],
            }
            constructed_data.append(split_data)

        if name in ["chemhivnoprompt"]:
            if name == "chemhivnoprompt":
                dataset = CHEMHIVOFADataset(
                    "chemhiv", sentence_encoder=encoder
                )
            # print(dataset[350342])
            # return 0
            split = dataset.get_idx_split()

            # def trim_class(embs, classes):
            #     valid_idx = classes == classes
            #     # valid_idx = torch.zeros_like(classes, dtype=torch.bool)
            #     return (
            #         torch.tensor([[0]]),
            #         embs[valid_idx.view(-1)].detach().clone(),
            #         classes[:, valid_idx.view(-1)].detach().clone(),
            #     )

            def trim_class(embs, label):
                label = label.to(torch.long)
                one_hot_label = torch.nn.functional.one_hot(
                    label, num_classes=2
                )
                return label, embs, one_hot_label

            # print(sub_dataset.label_text_feat.size())

            def make_data(split_name, state_name):
                return DataWithMeta(
                    GraphListNopromptDataset(
                        dataset,
                        dataset.label_text_feat,
                        dataset.prompt_edge_feat,
                        split[split_name],
                        single_prompt_edge=True,
                        # trim_class_func=trim_class,
                        walk_length=walk_length,
                    ),
                    batch_size,
                    sample_size=sample_size,
                    metric="auc",
                    state_name=state_name,
                    classes=len(dataset.label_text_feat),
                    meta_data={"eval_func": binary_auc_func},
                )

            split_data = {
                "train": [
                    GraphListNopromptDataset(
                        dataset,
                        dataset.label_text_feat,
                        dataset.prompt_edge_feat,
                        split["train"],
                        trim_class_func=trim_class,
                        single_prompt_edge=True,
                        walk_length=walk_length,
                    )
                ],
                "test": [
                    make_data("test", "test_" + name),
                    make_data("train", "test_train_" + name),
                ],
                "val": [make_data("valid", "valid_" + name)],
            }
            constructed_data.append(split_data)

        if name in ["coralink", "pubmedlink"]:
            if name == "coralink":
                dataset = CoraOFADataset("cora", sentence_encoder=encoder)
            else:
                dataset = PubmedOFADataset("pubmed", sentence_encoder=encoder)
            text_g = dataset.data
            text_g.x = text_g.x_text_feat
            text_g.prompt_edge_feat = dataset.prompt_edge_feat
            edges = text_g.edge_index
            edge_perm = torch.randperm(len(edges[0]))
            train_offset = int(len(edge_perm) * 0.85)
            val_offset = int(len(edge_perm) * 0.9)
            edge_indices = {
                "train": edge_perm[:train_offset],
                "val": edge_perm[train_offset:val_offset],
                "test": edge_perm[val_offset:],
            }
            graph_dict = text_g.to_dict()
            graph_dict["edge_index"] = edges[:, edge_indices["train"]]
            train_graph = pyg.data.Data(**graph_dict)

            def trim_class(label, num_class):
                binary_rep = torch.zeros((1, num_class))
                binary_rep[0, label] = 1
                return torch.tensor([label]).view(1, -1), binary_rep

            def make_data(split_name, state_name, remove_edge=False):
                return DataWithMeta(
                    SubgraphLinkHierDataset(
                        train_graph,
                        train_graph.edge_label_feat,
                        edges.T[edge_indices[split_name]].numpy(),
                        prompt_feat=train_graph.prompt_node_edge_feat,
                        to_undirected=True,
                        hop=3,
                        remove_edge=remove_edge,
                        trim_class_func=trim_class,
                        walk_length=walk_length,
                    ),
                    batch_size,
                    sample_size=sample_size,
                    metric="auc",
                    state_name=state_name,
                    classes=2,
                    meta_data={"eval_func": binary_auc_func},
                )

            split_data = {
                "train": [
                    SubgraphLinkHierDataset(
                        train_graph,
                        train_graph.edge_label_feat,
                        edges.T[edge_indices["train"]].numpy(),
                        prompt_feat=train_graph.prompt_node_edge_feat,
                        to_undirected=True,
                        hop=3,
                        remove_edge=True,
                        trim_class_func=trim_class,
                        walk_length=walk_length,
                    )
                ],
                "test": [
                    make_data("test", "test_" + name, remove_edge=False),
                    make_data("train", "test_train_" + name, remove_edge=True),
                ],
                "val": [make_data("val", "valid_" + name, remove_edge=False)],
            }
            constructed_data.append(split_data)

        if name in ["coranode", "pubmednode"]:
            if name == "coranode":
                dataset = CoraOFADataset("cora", sentence_encoder=encoder)
            else:
                dataset = PubmedOFADataset("pubmed", sentence_encoder=encoder)
            text_g = dataset.data
            text_g.x = text_g.x_text_feat
            text_g.prompt_edge_feat = dataset.prompt_edge_feat
            split = {
                "train": text_g.train_masks[0].nonzero(as_tuple=True)[0],
                "val": text_g.val_masks[0].nonzero(as_tuple=True)[0],
                "test": text_g.test_masks[0].nonzero(as_tuple=True)[0],
            }

            def trim_class(label, num_class):
                binary_rep = torch.zeros((1, num_class))
                binary_rep[0, label] = 1
                return torch.tensor([label]).view(1, -1), binary_rep

            def make_data(split_name, state_name):
                return DataWithMeta(
                    SubgraphHierDataset(
                        text_g,
                        text_g.label_text_feat,
                        split[split_name],
                        prompt_feat=text_g.prompt_node_feat,
                        to_undirected=True,
                        trim_class_func=trim_class,
                        walk_length=walk_length,
                    ),
                    batch_size,
                    sample_size=sample_size,
                    metric="acc",
                    state_name=state_name,
                    classes=len(text_g.label_text_feat),
                    meta_data={"eval_func": classification_func},
                )

            split_data = {
                "train": [
                    SubgraphHierDataset(
                        text_g,
                        text_g.label_text_feat,
                        split["train"],
                        prompt_feat=text_g.prompt_node_feat,
                        to_undirected=True,
                        trim_class_func=trim_class,
                        walk_length=walk_length,
                    )
                ],
                "test": [
                    make_data("test", "test_" + name),
                    make_data("train", "test_train_" + name),
                ],
                "val": [make_data("val", "valid_" + name)],
            }
            constructed_data.append(split_data)

        if name in ["WN18RR", "FB15K237"]:
            if name == "WN18RR":
                dataset = WN18RROFADataset("WN18RR", sentence_encoder=encoder)
            else:
                dataset = FB15K237OFADataset(
                    "FB15K237", sentence_encoder=encoder
                )
            text_g = dataset.data
            text_g.x = text_g.x_text_feat
            text_g.prompt_edge_feat = dataset.prompt_edge_feat
            edges = text_g.edge_index
            converted_triplet = dataset.get_idx_split()

            def trim_class(label, num_class):
                binary_rep = torch.zeros((1, num_class))
                binary_rep[0, label] = 1
                return torch.tensor([label]).view(1, -1), binary_rep

            def make_data(split_name, state_name, remove_edge=False):
                return DataWithMeta(
                    SubgraphKGHierDataset(
                        text_g,
                        text_g.edge_label_feat,
                        converted_triplet[split_name],
                        prompt_feat=text_g.prompt_text_feat,
                        to_undirected=True,
                        hop=2,
                        remove_edge=remove_edge,
                        trim_class_func=trim_class,
                        walk_length=walk_length,
                    ),
                    batch_size,
                    sample_size=sample_size,
                    metric="acc",
                    state_name=state_name,
                    classes=len(text_g.edge_label_feat),
                    meta_data={"eval_func": classification_func},
                )

            split_data = {
                "train": [
                    SubgraphKGHierDataset(
                        text_g,
                        text_g.edge_label_feat,
                        converted_triplet["train"],
                        prompt_feat=text_g.prompt_text_feat,
                        to_undirected=True,
                        hop=2,
                        remove_edge=True,
                        trim_class_func=trim_class,
                        walk_length=walk_length,
                    )
                ],
                "test": [
                    make_data("test", "test_" + name, remove_edge=False),
                    make_data("train", "test_train_" + name, remove_edge=True),
                ],
                "val": [
                    make_data("valid", "valid_" + name, remove_edge=False)
                ],
            }
            constructed_data.append(split_data)

        if name == "wikics":
            dataset = WikiCSOFADataset("wikics", sentence_encoder=encoder)
            text_g = dataset.data
            text_g.x = text_g.x_text_feat
            text_g.prompt_edge_feat = dataset.prompt_edge_feat
            # wikics provides 20 split combinations for train and val
            wiki_split_idx = 0
            text_split = [
                torch.where(text_g.train_mask[:, wiki_split_idx])[0].numpy(),
                torch.where(text_g.val_mask[:, wiki_split_idx])[0].numpy(),
                torch.where(text_g.test_mask)[0].numpy(),
            ]

            def trim_class(label, num_class):
                binary_rep = torch.zeros((1, num_class))
                binary_rep[0, label] = 1
                return label.view(1, -1), binary_rep

            def make_data(split_name, state_name):
                return DataWithMeta(
                    SubgraphHierDataset(
                        text_g,
                        text_g.label_text_feat,
                        text_split[split_name],
                        prompt_feat=text_g.prompt_text_feat,
                        to_undirected=False,
                        trim_class_func=trim_class,
                        walk_length=walk_length,
                    ),
                    batch_size,
                    sample_size=sample_size,
                    metric="acc",
                    state_name=state_name,
                    classes=10,
                    meta_data={"eval_func": classification_func},
                )

            split_data = {
                "train": [
                    SubgraphHierDataset(
                        text_g,
                        text_g.label_text_feat,
                        text_split[0],
                        prompt_feat=text_g.prompt_text_feat,
                        to_undirected=False,
                        trim_class_func=trim_class,
                        walk_length=walk_length,
                    )
                ],
                "test": [
                    make_data(2, "test_wikics"),
                    make_data(0, "test_train_wikics"),
                ],
                "val": [make_data(1, "valid_wikics")],
            }
            constructed_data.append(split_data)

        if name == "molzero":
            chembl_dataset = CHEMBLPREOFADataset(
                "chemblpre", sentence_encoder=encoder
            )
            hiv_dataset = CHEMHIVOFADataset(
                "chemhiv", sentence_encoder=encoder
            )
            pcba_dataset = CHEMPCBAOFADataset(
                "chempcba", sentence_encoder=encoder
            )
            hiv_split = hiv_dataset.get_idx_split()
            chembl_split = chembl_dataset.get_idx_split()
            pcba_split = pcba_dataset.get_idx_split()

            def trim_class(embs, classes):
                valid_idx = classes == classes
                # valid_idx = torch.zeros_like(classes, dtype=torch.bool)
                return (
                    torch.tensor([[0]]),
                    embs[valid_idx.view(-1)].detach().clone(),
                    classes[:, valid_idx.view(-1)].detach().clone(),
                )

            def hiv_trim_class(embs, label):
                # one_hot_label = torch.nn.functional.one_hot(
                #     label.to(torch.long), num_classes=2
                # )
                return label, embs[0:1], label

            def make_data(split_name, state_name):
                return DataWithMeta(
                    GraphListHierDataset(
                        hiv_dataset,
                        hiv_dataset.label_text_feat,
                        hiv_dataset.prompt_edge_feat,
                        hiv_dataset.prompt_text_feat,
                        hiv_split[split_name],
                        trim_class_func=hiv_trim_class,
                        single_prompt_edge=True,
                        walk_length=walk_length,
                    ),
                    batch_size,
                    sample_size=sample_size,
                    metric="auc",
                    state_name=state_name,
                    classes=1,
                    meta_data={"eval_func": binary_single_auc_func},
                )

            def make_pcba_data(split_name, state_name):
                return DataWithMeta(
                    GraphListHierDataset(
                        pcba_dataset,
                        pcba_dataset.label_text_feat,
                        pcba_dataset.prompt_edge_feat,
                        pcba_dataset.prompt_text_feat,
                        pcba_split[split_name],
                        # trim_class_func=trim_class,
                        single_prompt_edge=True,
                        walk_length=walk_length,
                    ),
                    batch_size,
                    sample_size=sample_size,
                    metric="aucmulti",
                    state_name=state_name,
                    classes=128,
                    meta_data={"eval_func": binary_auc_multi_func},
                )

            split_data = {
                "train": [
                    GraphListHierDataset(
                        chembl_dataset,
                        chembl_dataset.label_text_feat,
                        chembl_dataset.prompt_edge_feat,
                        chembl_dataset.prompt_text_feat,
                        chembl_split["train"],
                        trim_class_func=trim_class,
                        single_prompt_edge=True,
                        walk_length=walk_length,
                    )
                ],
                "test": [
                    make_data("test", "test_molhiv"),
                    # make_data("train", "test_train_molhiv"),
                    make_pcba_data("test", "test_pcba"),
                ],
                "val": [
                    make_data("valid", "valid_molhiv"),
                    make_pcba_data("valid", "valid_pcba"),
                ],
            }
            constructed_data.append(split_data)
        if name == "molfew":
            chembl_dataset = CHEMBLPREOFADataset(
                "chemblpre", sentence_encoder=encoder
            )
            hiv_dataset = CHEMHIVOFADataset(
                "chemhiv", sentence_encoder=encoder
            )
            pcba_dataset = CHEMPCBAOFADataset(
                "chempcba", sentence_encoder=encoder
            )
            pcba_split = pcba_dataset.get_idx_split()
            hiv_split = hiv_dataset.get_idx_split()
            chembl_split = chembl_dataset.get_idx_split()
            chembl_classes = chembl_dataset.y.view(len(chembl_dataset), -1)[
                chembl_split["train"]
            ]

            pcba_classes = [49, 60, 47, 94, 93]
            shots = [1, 3, 5, 10]

            def trim_class(embs, classes):
                valid_idx = classes == classes
                # valid_idx = torch.zeros_like(classes, dtype=torch.bool)
                return (
                    torch.tensor([[0]]),
                    embs[valid_idx.view(-1)].detach().clone(),
                    classes[:, valid_idx.view(-1)].detach().clone(),
                )

            def hiv_trim_class(embs, label):
                one_hot_label = torch.nn.functional.one_hot(
                    label.to(torch.long), num_classes=2
                )
                return label, embs, one_hot_label

            def make_data(split_name, state_name, shot=1):
                return DataWithMeta(
                    GraphListHierFSDataset(
                        hiv_dataset,
                        hiv_dataset.label_text_feat,
                        hiv_dataset.prompt_edge_feat,
                        hiv_dataset.prompt_text_feat,
                        hiv_split[split_name],
                        trim_class_func=hiv_trim_class,
                        single_prompt_edge=True,
                        walk_length=walk_length,
                        class_ind=hiv_dataset.y.view(len(hiv_dataset), -1)[
                            hiv_split[split_name], 0:1
                        ],
                        shot=shot,
                    ),
                    batch_size,
                    sample_size=sample_size,
                    metric="auc",
                    state_name=state_name,
                    classes=1,
                    meta_data={"eval_func": binary_single_auc_func},
                )

            def make_pcba_data(
                split_name, state_name, shot=1, target_class=None
            ):
                return DataWithMeta(
                    GraphListHierFSDataset(
                        pcba_dataset,
                        pcba_dataset.label_text_feat,
                        pcba_dataset.prompt_edge_feat,
                        pcba_dataset.prompt_text_feat,
                        pcba_split[split_name],
                        single_prompt_edge=True,
                        walk_length=walk_length,
                        class_ind=pcba_dataset.y.view(len(pcba_dataset), -1)[
                            pcba_split[split_name]
                        ],
                        shot=shot,
                        target_class=target_class,
                    ),
                    batch_size,
                    sample_size=sample_size,
                    metric="auc",
                    state_name=state_name,
                    classes=1,
                    meta_data={"eval_func": binary_single_auc_func},
                )

            # def hiv_trim_class(embs, label):
            #     return label.view(1, -1), embs[0:1], label.view(1, -1)

            # def make_data(split_name, state_name, shot=5):
            #     return DataWithMeta(
            #         GraphListHierFixDataset(
            #             hiv_dataset,
            #             hiv_dataset.label_text_feat,
            #             hiv_dataset.prompt_edge_feat,
            #             hiv_dataset.prompt_text_feat,
            #             hiv_split[split_name],
            #             trim_class_func=hiv_trim_class,
            #             single_prompt_edge=True,
            #             walk_length=walk_length,
            #             class_ind=hiv_dataset.y.view(len(hiv_dataset), -1)[
            #                 hiv_split[split_name], 0:1
            #             ],
            #             shot=shot,
            #         ),
            #         batch_size,
            #         sample_size=sample_size,
            #         metric="auc",
            #         state_name=state_name,
            #         classes=1,
            #         meta_data={"eval_func": binary_single_auc_func},
            #     )
            hiv_val_data = [
                make_data("valid", "valid_hiv_" + str(i), i) for i in shots
            ]
            hiv_test_data = [
                make_data("test", "test_hiv_" + str(i), i) for i in shots
            ]

            pcba_val_data = [
                make_pcba_data(
                    "valid",
                    "valid_pcba_" + str(i) + "_" + str(j),
                    i,
                    torch.tensor([j]),
                )
                for i in shots
                for j in pcba_classes
            ]
            pcba_test_data = [
                make_pcba_data(
                    "test",
                    "test_pcba_" + str(i) + "_" + str(j),
                    i,
                    torch.tensor([j]),
                )
                for i in shots
                for j in pcba_classes
            ]

            split_data = {
                "train": [
                    GraphListHierFSDataset(
                        chembl_dataset,
                        chembl_dataset.label_text_feat,
                        chembl_dataset.prompt_edge_feat,
                        chembl_dataset.prompt_text_feat,
                        chembl_split["train"],
                        trim_class_func=trim_class,
                        single_prompt_edge=True,
                        walk_length=walk_length,
                        class_ind=chembl_classes,
                    )
                ],
                "test": hiv_test_data + pcba_test_data,
                "val": hiv_val_data + pcba_val_data,
            }
            constructed_data.append(split_data)
        if name == "molfewfull":
            chembl_dataset = CHEMBLPREOFADataset(
                "chemblpre", sentence_encoder=encoder
            )
            hiv_dataset = CHEMHIVOFADataset(
                "chemhiv", sentence_encoder=encoder
            )
            pcba_dataset = CHEMPCBAOFADataset(
                "chempcba", sentence_encoder=encoder
            )
            pcba_split = pcba_dataset.get_idx_split()
            hiv_split = hiv_dataset.get_idx_split()
            chembl_split = chembl_dataset.get_idx_split()
            chembl_classes = chembl_dataset.y.view(len(chembl_dataset), -1)[
                chembl_split["train"]
            ]

            def trim_class(embs, classes):
                valid_idx = classes == classes
                # valid_idx = torch.zeros_like(classes, dtype=torch.bool)
                return (
                    torch.tensor([[0]]),
                    embs[valid_idx.view(-1)].detach().clone(),
                    classes[:, valid_idx.view(-1)].detach().clone(),
                )

            def hiv_trim_class(embs, label):
                return label.view(1, -1), embs[0:1], label.view(1, -1)

            def make_data(split_name, state_name, shot=5):
                return DataWithMeta(
                    GraphListHierFixDataset(
                        hiv_dataset,
                        hiv_dataset.label_text_feat,
                        hiv_dataset.prompt_edge_feat,
                        hiv_dataset.prompt_text_feat,
                        hiv_split[split_name],
                        trim_class_func=hiv_trim_class,
                        single_prompt_edge=True,
                        walk_length=walk_length,
                        class_ind=hiv_dataset.y.view(len(hiv_dataset), -1)[
                            hiv_split[split_name], 0:1
                        ],
                        shot=1,
                    ),
                    batch_size,
                    sample_size=sample_size,
                    metric="auc",
                    state_name=state_name,
                    classes=2,
                    meta_data={"eval_func": binary_single_auc_func},
                )

            def make_pcba_data(split_name, state_name, shot=5):
                return DataWithMeta(
                    GraphListHierFixDataset(
                        pcba_dataset,
                        pcba_dataset.label_text_feat,
                        hiv_dataset.prompt_edge_feat,
                        pcba_dataset.prompt_text_feat,
                        pcba_split[split_name],
                        trim_class_func=None,
                        single_prompt_edge=True,
                        walk_length=walk_length,
                        class_ind=pcba_dataset.y.view(len(pcba_dataset), -1)[
                            pcba_split[split_name]
                        ],
                        shot=5,
                    ),
                    batch_size,
                    sample_size=sample_size,
                    metric="apr",
                    state_name=state_name,
                    classes=len(pcba_dataset.label_text_feat),
                    meta_data={"eval_func": binary_apr_func},
                )

            split_data = {
                "train": [
                    GraphListHierFSDataset(
                        chembl_dataset,
                        chembl_dataset.label_text_feat,
                        chembl_dataset.prompt_edge_feat,
                        chembl_dataset.prompt_text_feat,
                        chembl_split["train"],
                        trim_class_func=trim_class,
                        single_prompt_edge=True,
                        walk_length=walk_length,
                        class_ind=chembl_classes,
                        shot=2,
                    )
                ],
                "test": [
                    make_data("test", "test_molhiv"),
                    # make_data("train", "test_train_molhiv"),
                    # make_pcba_data("test", "test_pcba"),
                    # make_data("train", "test_train_molhiv"),
                ],
                "val": [
                    make_data("valid", "valid_molhiv"),
                    # make_pcba_data("valid", "valid_pcba"),
                ],
            }
            constructed_data.append(split_data)
    states = ["train", "test", "val"]
    collate_dataset = {}
    for s in states:
        dataset_list = []
        for ds in constructed_data:
            dataset_list += ds[s]
        collate_dataset[s] = dataset_list
    return collate_dataset


def main(params):
    encoder = SentenceEncoder("ST")

    if hasattr(params, "data_list"):
        e2e_data_list = params.data_list.split(",")
    else:
        e2e_data_list = [
            "coralink",
            "coranode",
            "pubmedlink",
            "pubmednode",
            "arxiv",
            "WN18RR",
            "FB15K237",
            "wikics",
            "chemblpre",
            "chempcba",
            "chemhiv",
        ]

    collate_dataset = data_construct(
        e2e_data_list,
        encoder,
        batch_size=params.eval_batch_size,
        sample_size=params.eval_sample_size,
        walk_length=params.rwpe,
    )

    out_dim = 768 + (params.rwpe if params.rwpe is not None else 0)
    # out_dim = 768

    def make_data(data, b_size, sample_size):
        return DataWithMeta(data, b_size, sample_size=sample_size)

    if hasattr(params, "d_multiple"):
        data_multiple = [float(a) for a in params.d_multiple.split(",")]
    else:
        data_multiple = [2, 2, 0.3, 2, 0.5, 0.4, 0.3, 2, 1, 2, 3]

    if hasattr(params, "d_min_ratio"):
        min_ratio = [float(a) for a in params.d_min_ratio.split(",")]
    else:
        min_ratio = [0.5, 0.5, 0.05, 1, 0.1, 0.1, 0.03, 1, 0.2, 0.2, 1]

    train_data = MultiDataset(
        collate_dataset["train"],
        dataset_multiple=data_multiple,
        # dataset_multiple=1,
        patience=3,
        window_size=5,
        min_ratio=min_ratio,
        # min_ratio=1,
    )

    text_dataset = {
        "train": make_data(
            train_data,
            params.batch_size,
            params.train_sample_size,
        ),
        "val": collate_dataset["val"],
        "test": collate_dataset["test"],
    }
    params.datamodule = DataModule(
        text_dataset, num_workers=params.num_workers
    )

    eval_data = text_dataset["val"] + text_dataset["test"]
    val_state = [dt.state_name for dt in text_dataset["val"]]
    test_state = [dt.state_name for dt in text_dataset["test"]]
    eval_state = val_state + test_state
    eval_metric = [dt.metric for dt in eval_data]
    eval_funcs = [dt.meta_data["eval_func"] for dt in eval_data]
    loss = torch.nn.BCEWithLogitsLoss()
    evlter = []
    for dt in eval_data:
        if dt.metric == "acc":
            evlter.append(Accuracy(task="multiclass", num_classes=dt.classes))
        elif dt.metric == "auc":
            evlter.append(AUROC(task="binary"))
        elif dt.metric == "apr":
            evlter.append(MultiApr(num_labels=dt.classes))
        elif dt.metric == "aucmulti":
            evlter.append(MultiAuc(num_labels=dt.classes))
    metrics = EvalKit(
        eval_metric,
        evlter,
        loss,
        eval_funcs,
        flat_binary_func,
        eval_mode="max",
        exp_prefix="",
        eval_state=eval_state,
        val_monitor_state=val_state[0],
        test_monitor_state=test_state[0],
    )
    # gnn = PyGGIN(params.num_layers, 768, 768)
    # gnn = PyGRGCN(params.num_layers, 3, 768, 768)
    # gnn = PyGGINE(params.num_layers, 768, 768, 768)
    gnn = PyGRGCNEdge(
        params.num_layers,
        5,
        out_dim,
        out_dim,
        drop_ratio=params.dropout,
        JK=params.JK,
    )
    bin_model = BinGraphAttModel if params.JK == "none" else BinGraphModel
    model = bin_model(gnn, out_dim, 1, add_rwpe=params.rwpe)
    optimizer = torch.optim.Adam(
        model.parameters(), lr=params.lr, weight_decay=params.l2
    )
    lr_scheduler = {
        "scheduler": torch.optim.lr_scheduler.StepLR(optimizer, 15, 0.5),
        "interval": "epoch",
        "frequency": 1,
    }

    exp_config = ExpConfig(
        "",
        optimizer,
        dataset_callback=train_data.update,
        lr_scheduler=lr_scheduler,
    )
    exp_config.val_state_name = val_state
    exp_config.test_state_name = test_state

    pred_model = GraphPredLightning(exp_config, model, metrics)

    wandb_logger = WandbLogger(
        project=params.log_project,
        name=params.exp_name,
        save_dir=params.exp_dir,
        offline=params.offline_log,
    )

    val_res, test_res = lightning_fit(
        wandb_logger,
        pred_model,
        params.datamodule,
        metrics,
        params.num_epochs,
        save_model=False,
        load_best=False,
        reload_freq=1,
        test_rep=params.test_rep
        # profiler="simple",
        # accelerator="cpu",
    )


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="rl")

    parser.add_argument(
        "opts",
        default=[],
        nargs=argparse.REMAINDER,
        help="Modify config options using the command-line",
    )

    parser.add_argument("--n_way", type=int, default=5)
    parser.add_argument("--k_shot", type=int, default=3)
    parser.add_argument("--q_query", type=int, default=3)
    parser.add_argument(
        "--fs_task_num",
        type=int,
        default=5,
        help="Number of tasks for few-shot training.",
    )

    params = parser.parse_args()
    configs = []
    configs.append(
        load_yaml(
            os.path.join(
                os.path.dirname(__file__), "configs", "default_config.yaml"
            )
        )
    )
    print(configs)
    # Add for few-shot parameters
    configs.append(params.__dict__)

    mod_params = combine_dict(*configs)
    mod_params = merge_mod(mod_params, params.opts)
    setup_exp(mod_params)

    params = SimpleNamespace(**mod_params)
    set_random_seed(params.seed)

    torch.set_float32_matmul_precision("high")
    params.log_project = "full_cdm"
    main(params)

run_cdm.py
LechengKong/OneForAll
import argparse
import os
import torch
import torch_geometric as pyg
from pytorch_lightning.loggers import WandbLogger

from gp.utils.utils import (
    load_yaml,
    combine_dict,
    merge_mod,
    setup_exp,
    set_random_seed,
    k_fold2_split,
    k_fold_ind,
)
from gp.lightning.metric import (
    binary_auc_func,
    flat_binary_func,
    classification_func,
    EvalKit,
)
from gp.lightning.data_template import DataModule, DataWithMeta
from gp.lightning.training import lightning_fit, lightning_test
from gp.lightning.module_template import ExpConfig
from types import SimpleNamespace
from lightning_model import GraphPredLightning
from models.model import BinGraphModel, BinGraphAttModel, AdaPoolClassModel
from gp.nn.models.pyg import PyGGIN, PyGRGCN, PyGGINE
from models.model import PyGRGCNEdge

from torchmetrics import AUROC, Accuracy
from utils import (
    SentenceEncoder,
    binary_apr_func,
    MultiApr,
    flat_label_func,
    label_apr_func,
    binary_single_auc_func,
    binary_auc_multi_func,
    MultiAuc,
)
from gp.lightning.metric import flat_binary_func_fs
from gp.utils.utils import SmartTimer
from scipy.sparse import csr_array

from task_constructor import TaskConstructor


def main(params):
    encoder = SentenceEncoder("ST")

    tasks = TaskConstructor(["chemhiv"], encoder)

    out_dim = 768 + (params.rwpe if params.rwpe is not None else 0)
    # out_dim = 768

    if hasattr(params, "d_multiple"):
        if isinstance(params.d_multiple, str):
            data_multiple = [float(a) for a in params.d_multiple.split(",")]
        else:
            data_multiple = params.d_multiple
    else:
        data_multiple = [2, 2, 0.3, 2, 0.5, 0.4, 0.3, 2, 1, 2, 3]

    if hasattr(params, "d_min_ratio"):
        if isinstance(params.d_min_ratio, str):
            min_ratio = [float(a) for a in params.d_min_ratio.split(",")]
        else:
            min_ratio = params.d_min_ratio
    else:
        min_ratio = [0.5, 0.5, 0.05, 1, 0.1, 0.1, 0.03, 1, 0.2, 0.2, 1]

    train_data = tasks.make_train_data(data_multiple, min_ratio)

    text_dataset = tasks.make_full_dm_list(
        data_multiple, min_ratio, train_data
    )
    params.datamodule = DataModule(
        text_dataset, num_workers=params.num_workers
    )

    eval_data = text_dataset["val"] + text_dataset["test"]
    val_state = [dt.state_name for dt in text_dataset["val"]]
    test_state = [dt.state_name for dt in text_dataset["test"]]
    eval_state = val_state + test_state
    eval_metric = [dt.metric for dt in eval_data]
    eval_funcs = [dt.meta_data["eval_func"] for dt in eval_data]
    loss = torch.nn.BCEWithLogitsLoss()
    evlter = []
    for dt in eval_data:
        if dt.metric == "acc":
            evlter.append(Accuracy(task="multiclass", num_classes=dt.classes))
        elif dt.metric == "auc":
            evlter.append(AUROC(task="binary"))
        elif dt.metric == "apr":
            evlter.append(MultiApr(num_labels=dt.classes))
        elif dt.metric == "aucmulti":
            evlter.append(MultiAuc(num_labels=dt.classes))
    metrics = EvalKit(
        eval_metric,
        evlter,
        loss,
        eval_funcs,
        flat_binary_func,
        eval_mode="max",
        exp_prefix="",
        eval_state=eval_state,
        val_monitor_state=val_state[0],
        test_monitor_state=test_state[0],
    )
    # gnn = PyGGIN(params.num_layers, 768, 768)
    # gnn = PyGRGCN(params.num_layers, 3, 768, 768)
    # gnn = PyGGINE(params.num_layers, 768, 768, 768)
    gnn = PyGRGCNEdge(
        params.num_layers,
        5,
        out_dim,
        out_dim,
        drop_ratio=params.dropout,
        JK=params.JK,
    )
    bin_model = BinGraphAttModel if params.JK == "none" else BinGraphModel
    model = bin_model(gnn, out_dim, 1, add_rwpe=params.rwpe)
    optimizer = torch.optim.Adam(
        model.parameters(), lr=params.lr, weight_decay=params.l2
    )
    lr_scheduler = {
        "scheduler": torch.optim.lr_scheduler.StepLR(optimizer, 15, 0.5),
        "interval": "epoch",
        "frequency": 1,
    }

    exp_config = ExpConfig(
        "",
        optimizer,
        dataset_callback=train_data.update,
        lr_scheduler=lr_scheduler,
    )
    exp_config.val_state_name = val_state
    exp_config.test_state_name = test_state

    pred_model = GraphPredLightning(exp_config, model, metrics)

    wandb_logger = WandbLogger(
        project=params.log_project,
        name=params.exp_name,
        save_dir=params.exp_dir,
        offline=params.offline_log,
    )

    val_res, test_res = lightning_fit(
        wandb_logger,
        pred_model,
        params.datamodule,
        metrics,
        params.num_epochs,
        save_model=False,
        load_best=False,
        reload_freq=1,
        test_rep=params.test_rep
        # profiler="simple",
        # accelerator="cpu",
    )


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="rl")

    parser.add_argument(
        "opts",
        default=[],
        nargs=argparse.REMAINDER,
        help="Modify config options using the command-line",
    )

    parser.add_argument("--n_way", type=int, default=5)
    parser.add_argument("--k_shot", type=int, default=3)
    parser.add_argument("--q_query", type=int, default=3)
    parser.add_argument(
        "--fs_task_num",
        type=int,
        default=5,
        help="Number of tasks for few-shot training.",
    )

    params = parser.parse_args()
    configs = []
    configs.append(
        load_yaml(
            os.path.join(
                os.path.dirname(__file__), "configs", "default_config.yaml"
            )
        )
    )
    print(configs)
    # Add for few-shot parameters
    configs.append(params.__dict__)

    mod_params = combine_dict(*configs)
    mod_params = merge_mod(mod_params, params.opts)
    setup_exp(mod_params)

    params = SimpleNamespace(**mod_params)
    set_random_seed(params.seed)

    torch.set_float32_matmul_precision("high")
    params.log_project = "full_cdm"
    main(params)

run_cdm_ele.py
LechengKong/OneForAll
import argparse
import os
import torch
import torch_geometric as pyg
from pytorch_lightning.loggers import WandbLogger
from ofa_datasets import (
    GraphListDataset,
    SubgraphDataset,
    MultiDataset,
    GraphListHierDataset,
    SubgraphHierDataset,
    GraphListHierFSDataset,
    GraphListHierFixDataset,
    SubgraphLinkHierDataset,
    SubgraphKGHierDataset,
    FewShotNCDataset,
    FewShotKGDataset,
    ZeroShotNCDataset,
    ZeroShotKGDataset,
)
from fs_datamanager import FewShotDataManager
from gp.utils.utils import (
    load_yaml,
    combine_dict,
    merge_mod,
    setup_exp,
    set_random_seed,
    k_fold2_split,
    k_fold_ind,
)
from gp.lightning.metric import (
    binary_auc_func,
    flat_binary_func,
    classification_func,
    EvalKit,
)
from gp.lightning.data_template import DataModule, DataWithMeta
from gp.lightning.training import lightning_fit, lightning_test
from gp.lightning.module_template import ExpConfig
from types import SimpleNamespace
from lightning_model import GraphPredLightning
from models.model import BinGraphModel, BinGraphAttModel
from gp.nn.models.pyg import PyGGIN, PyGRGCN, PyGGINE
from models.model import PyGRGCNEdge

from torchmetrics import AUROC, Accuracy
from utils import (
    SentenceEncoder,
    binary_apr_func,
    MultiApr,
    flat_label_func,
    label_apr_func,
    binary_single_auc_func,
    binary_auc_multi_func,
    MultiAuc,
    classification_single_func,
)
from data.arxiv.gen_data import ArxivOFADataset
from data.Cora.gen_data import CoraOFADataset
from data.WN18RR.gen_data import WN18RROFADataset
from data.FB15K237.gen_data import FB15K237OFADataset
from data.chemblpre.gen_data import CHEMBLPREOFADataset
from data.chempcba.gen_data import CHEMPCBAOFADataset
from data.chemhiv.gen_data import CHEMHIVOFADataset

from gp.lightning.metric import flat_binary_func_fs
from gp.utils.utils import SmartTimer
from scipy.sparse import csr_array


def data_construct(
    data_names,
    encoder,
    batch_size=None,
    sample_size=None,
    walk_length=None,
    fs_setting=None,
):
    constructed_data = []
    wiki_train_g = None
    for data_idx, name in enumerate(data_names):
        print("loading: ", name)
        if name in ["chempcba", "chemblpre"]:
            if name == "chempcba":
                dataset = CHEMPCBAOFADataset(
                    "chempcba", sentence_encoder=encoder
                )
            elif name == "chemblpre":
                dataset = CHEMBLPREOFADataset(
                    "chemblpre", sentence_encoder=encoder
                )
            # print(dataset[350342])
            # return 0
            split = dataset.get_idx_split()

            def trim_class(embs, classes):
                valid_idx = classes == classes
                # valid_idx = torch.zeros_like(classes, dtype=torch.bool)
                return (
                    torch.tensor([[0]]),
                    embs[valid_idx.view(-1)].detach().clone(),
                    classes[:, valid_idx.view(-1)].detach().clone(),
                )

            def make_data(split_name, state_name):
                return DataWithMeta(
                    GraphListHierDataset(
                        dataset,
                        dataset.label_text_feat,
                        dataset.prompt_edge_feat,
                        dataset.prompt_text_feat,
                        split[split_name],
                        single_prompt_edge=True,
                        # trim_class_func=trim_class,
                        walk_length=walk_length,
                    ),
                    batch_size,
                    sample_size=sample_size,
                    metric="apr",
                    state_name=state_name,
                    classes=len(dataset.label_text_feat),
                    meta_data={"eval_func": binary_apr_func},
                )

            split_data = {
                "train": [
                    GraphListHierDataset(
                        dataset,
                        dataset.label_text_feat,
                        dataset.prompt_edge_feat,
                        dataset.prompt_text_feat,
                        split["train"],
                        trim_class_func=trim_class,
                        single_prompt_edge=True,
                        walk_length=walk_length,
                    )
                ],
                "test": [
                    make_data("test", "test_" + name),
                    make_data("train", "test_train_" + name),
                ],
                "val": [make_data("valid", "valid_" + name)],
            }
            constructed_data.append(split_data)
        if name in ["chemhiv"]:
            if name == "chemhiv":
                dataset = CHEMHIVOFADataset(
                    "chemhiv", sentence_encoder=encoder
                )
            # print(dataset[350342])
            # return 0
            split = dataset.get_idx_split()

            # def trim_class(embs, classes):
            #     valid_idx = classes == classes
            #     # valid_idx = torch.zeros_like(classes, dtype=torch.bool)
            #     return (
            #         torch.tensor([[0]]),
            #         embs[valid_idx.view(-1)].detach().clone(),
            #         classes[:, valid_idx.view(-1)].detach().clone(),
            #     )

            def trim_class(embs, label):
                label = label.to(torch.long)
                one_hot_label = torch.nn.functional.one_hot(
                    label, num_classes=2
                )
                return label, embs, one_hot_label

            # print(sub_dataset.label_text_feat.size())

            def make_data(split_name, state_name):
                return DataWithMeta(
                    GraphListHierDataset(
                        dataset,
                        dataset.label_text_feat,
                        dataset.prompt_edge_feat,
                        dataset.prompt_text_feat,
                        split[split_name],
                        single_prompt_edge=True,
                        # trim_class_func=trim_class,
                        walk_length=walk_length,
                    ),
                    batch_size,
                    sample_size=sample_size,
                    metric="auc",
                    state_name=state_name,
                    classes=len(dataset.label_text_feat),
                    meta_data={"eval_func": binary_auc_func},
                )

            split_data = {
                "train": [
                    GraphListHierDataset(
                        dataset,
                        dataset.label_text_feat,
                        dataset.prompt_edge_feat,
                        dataset.prompt_text_feat,
                        split["train"],
                        trim_class_func=trim_class,
                        single_prompt_edge=True,
                        walk_length=walk_length,
                    )
                ],
                "test": [
                    make_data("test", "test_" + name),
                    make_data("train", "test_train_" + name),
                ],
                "val": [make_data("valid", "valid_" + name)],
            }
            constructed_data.append(split_data)

        if name == "molzero":
            chembl_dataset = CHEMBLPREOFADataset(
                "chemblpre", sentence_encoder=encoder
            )
            hiv_dataset = CHEMHIVOFADataset(
                "chemhiv", sentence_encoder=encoder
            )
            pcba_dataset = CHEMPCBAOFADataset(
                "chempcba", sentence_encoder=encoder
            )
            hiv_split = hiv_dataset.get_idx_split()
            chembl_split = chembl_dataset.get_idx_split()
            pcba_split = pcba_dataset.get_idx_split()

            def trim_class(embs, classes):
                valid_idx = classes == classes
                # valid_idx = torch.zeros_like(classes, dtype=torch.bool)
                return (
                    torch.tensor([[0]]),
                    embs[valid_idx.view(-1)].detach().clone(),
                    classes[:, valid_idx.view(-1)].detach().clone(),
                )

            def hiv_trim_class(embs, label):
                # one_hot_label = torch.nn.functional.one_hot(
                #     label.to(torch.long), num_classes=2
                # )
                return label, embs[0:1], label

            def make_data(split_name, state_name):
                return DataWithMeta(
                    GraphListHierDataset(
                        hiv_dataset,
                        hiv_dataset.label_text_feat,
                        hiv_dataset.prompt_edge_feat,
                        hiv_dataset.prompt_text_feat,
                        hiv_split[split_name],
                        trim_class_func=hiv_trim_class,
                        single_prompt_edge=True,
                        walk_length=walk_length,
                    ),
                    batch_size,
                    sample_size=sample_size,
                    metric="auc",
                    state_name=state_name,
                    classes=1,
                    meta_data={"eval_func": binary_single_auc_func},
                )

            def make_pcba_data(split_name, state_name):
                return DataWithMeta(
                    GraphListHierDataset(
                        pcba_dataset,
                        pcba_dataset.label_text_feat,
                        pcba_dataset.prompt_edge_feat,
                        pcba_dataset.prompt_text_feat,
                        pcba_split[split_name],
                        # trim_class_func=trim_class,
                        single_prompt_edge=True,
                        walk_length=walk_length,
                    ),
                    batch_size,
                    sample_size=sample_size,
                    metric="aucmulti",
                    state_name=state_name,
                    classes=128,
                    meta_data={"eval_func": binary_auc_multi_func},
                )

            split_data = {
                "train": [
                    GraphListHierDataset(
                        chembl_dataset,
                        chembl_dataset.label_text_feat,
                        chembl_dataset.prompt_edge_feat,
                        chembl_dataset.prompt_text_feat,
                        chembl_split["train"],
                        trim_class_func=trim_class,
                        single_prompt_edge=True,
                        walk_length=walk_length,
                    )
                ],
                "test": [
                    make_data("test", "test_molhiv"),
                    # make_data("train", "test_train_molhiv"),
                    make_pcba_data("test", "test_pcba"),
                ],
                "val": [
                    make_data("valid", "valid_molhiv"),
                    make_pcba_data("valid", "valid_pcba"),
                ],
            }
            constructed_data.append(split_data)
        if name == "molfew":
            chembl_dataset = CHEMBLPREOFADataset(
                "chemblpre", sentence_encoder=encoder
            )
            hiv_dataset = CHEMHIVOFADataset(
                "chemhiv", sentence_encoder=encoder
            )
            pcba_dataset = CHEMPCBAOFADataset(
                "chempcba", sentence_encoder=encoder
            )
            pcba_split = pcba_dataset.get_idx_split()
            hiv_split = hiv_dataset.get_idx_split()
            chembl_split = chembl_dataset.get_idx_split()
            chembl_classes = chembl_dataset.y.view(len(chembl_dataset), -1)[
                chembl_split["train"]
            ]

            pcba_classes = [49, 60, 47, 94, 93]
            shots = [1, 3, 5, 10]

            def trim_class(embs, classes):
                valid_idx = classes == classes
                # valid_idx = torch.zeros_like(classes, dtype=torch.bool)
                return (
                    torch.tensor([[0]]),
                    embs[valid_idx.view(-1)].detach().clone(),
                    classes[:, valid_idx.view(-1)].detach().clone(),
                )

            def hiv_trim_class(embs, label):
                one_hot_label = torch.nn.functional.one_hot(
                    label.to(torch.long), num_classes=2
                )
                return label, embs, one_hot_label

            def make_data(split_name, state_name, shot=1):
                return DataWithMeta(
                    GraphListHierFSDataset(
                        hiv_dataset,
                        hiv_dataset.label_text_feat,
                        hiv_dataset.prompt_edge_feat,
                        hiv_dataset.prompt_text_feat,
                        hiv_split[split_name],
                        trim_class_func=hiv_trim_class,
                        single_prompt_edge=True,
                        walk_length=walk_length,
                        class_ind=hiv_dataset.y.view(len(hiv_dataset), -1)[
                            hiv_split[split_name], 0:1
                        ],
                        shot=shot,
                    ),
                    batch_size,
                    sample_size=1000,
                    metric="auc",
                    state_name=state_name,
                    classes=1,
                    meta_data={"eval_func": binary_single_auc_func},
                )

            def make_pcba_data(
                split_name, state_name, shot=1, target_class=None
            ):
                return DataWithMeta(
                    GraphListHierFSDataset(
                        pcba_dataset,
                        pcba_dataset.label_text_feat,
                        pcba_dataset.prompt_edge_feat,
                        pcba_dataset.prompt_text_feat,
                        pcba_split[split_name],
                        single_prompt_edge=True,
                        walk_length=walk_length,
                        class_ind=pcba_dataset.y.view(len(pcba_dataset), -1)[
                            pcba_split[split_name]
                        ],
                        shot=shot,
                        target_class=target_class,
                    ),
                    batch_size,
                    sample_size=1000,
                    metric="auc",
                    state_name=state_name,
                    classes=1,
                    meta_data={"eval_func": binary_single_auc_func},
                )

            # def hiv_trim_class(embs, label):
            #     return label.view(1, -1), embs[0:1], label.view(1, -1)

            # def make_data(split_name, state_name, shot=5):
            #     return DataWithMeta(
            #         GraphListHierFixDataset(
            #             hiv_dataset,
            #             hiv_dataset.label_text_feat,
            #             hiv_dataset.prompt_edge_feat,
            #             hiv_dataset.prompt_text_feat,
            #             hiv_split[split_name],
            #             trim_class_func=hiv_trim_class,
            #             single_prompt_edge=True,
            #             walk_length=walk_length,
            #             class_ind=hiv_dataset.y.view(len(hiv_dataset), -1)[
            #                 hiv_split[split_name], 0:1
            #             ],
            #             shot=shot,
            #         ),
            #         batch_size,
            #         sample_size=sample_size,
            #         metric="auc",
            #         state_name=state_name,
            #         classes=1,
            #         meta_data={"eval_func": binary_single_auc_func},
            #     )
            hiv_val_data = [
                make_data("valid", "valid_hiv_" + str(i), i) for i in shots
            ]
            hiv_test_data = [
                make_data("test", "test_hiv_" + str(i), i) for i in shots
            ]

            pcba_val_data = [
                make_pcba_data(
                    "valid",
                    "valid_pcba_" + str(i) + "_" + str(j),
                    i,
                    torch.tensor([j]),
                )
                for i in shots
                for j in pcba_classes
            ]
            pcba_test_data = [
                make_pcba_data(
                    "test",
                    "test_pcba_" + str(i) + "_" + str(j),
                    i,
                    torch.tensor([j]),
                )
                for i in shots
                for j in pcba_classes
            ]

            split_data = {
                "train": [
                    GraphListHierFSDataset(
                        chembl_dataset,
                        chembl_dataset.label_text_feat,
                        chembl_dataset.prompt_edge_feat,
                        chembl_dataset.prompt_text_feat,
                        chembl_split["train"],
                        trim_class_func=trim_class,
                        single_prompt_edge=True,
                        walk_length=walk_length,
                        class_ind=chembl_classes,
                    )
                ],
                "test": hiv_test_data + pcba_test_data,
                "val": hiv_val_data + pcba_val_data,
            }
            constructed_data.append(split_data)

        if name.startswith("fs_"):

            (
                [n_way, min_n],
                [k_shot, min_k],
                q_query,
                fs_task_num,
                class_emb_flag,
            ) = fs_setting[name]

            def make_data(
                g,
                data_manager,
                split_name: int,
                state_name,
                undirected_flag,
                adj=None,
                class_emb_flag=False,
                train_flag=False,
                random_flag=False,
                min_n=None,
                min_k=None,
            ):
                if class_emb_flag:
                    class_emb = g.label_text_feat
                else:
                    class_emb = g.prompt_text_feat.repeat(
                        len(g.label_text_feat), 1
                    )
                if k_shot > 0:
                    data_class = FewShotNCDataset
                else:
                    data_class = ZeroShotNCDataset
                dataset = data_class(
                    pyg_graph=g,
                    class_emb=class_emb,
                    data_idx=torch.zeros(50),
                    n_way=n_way,
                    k_shot=k_shot,
                    q_query=q_query,
                    datamanager=data_manager,
                    mode=split_name,
                    hop=2,
                    prompt_feat=g.prompt_text_feat,
                    to_undirected=undirected_flag,
                    adj=adj,
                    single_prompt_edge=True,
                    random_flag=random_flag,
                    min_n=min_n,
                    min_k=min_k,
                )
                if train_flag:
                    return dataset
                return DataWithMeta(
                    dataset,
                    batch_size=fs_task_num,
                    sample_size=-1,
                    metric="acc",
                    state_name=state_name,
                    classes=min_n,
                    meta_data={"eval_func": classification_single_func},
                )

            if "arxiv" in name:
                dataset = ArxivOFADataset("arxiv", sentence_encoder=encoder)
                text_g = dataset.data
                text_g.x = text_g.x_text_feat
                text_g.prompt_edge_feat = dataset.prompt_edge_feat
                data_manager = FewShotDataManager(
                    text_g,
                    n_way,
                    k_shot,
                    q_query,
                )
                # only use cora for test
                dataset_cora = CoraOFADataset("cora", sentence_encoder=encoder)
                g_cora = dataset_cora.data
                g_cora.x = dataset_cora.x_text_feat
                g_cora.prompt_edge_feat = dataset_cora.prompt_edge_feat
                g_cora.prompt_text_feat = g_cora.prompt_node_feat
                cora_fsdm = FewShotDataManager(
                    g_cora, n_way, k_shot, q_query, class_split_ratio=[7, 0, 0]
                )

                if k_shot > 0:
                    k_shot_lst = [1, 3, 5]
                else:
                    k_shot_lst = [0]

                train_data = [
                    make_data(
                        text_g,
                        data_manager,
                        0,
                        f"train_fs_arxiv",
                        True,
                        class_emb_flag=class_emb_flag,
                        train_flag=True,
                        random_flag=True,
                        min_n=min_n,
                        min_k=min_k,
                    )
                ]
                valid_data = [
                    make_data(
                        g[0],
                        g[1],
                        g[2],
                        f"val_fs{n}{k}_{g[3]}",
                        True,
                        class_emb_flag=class_emb_flag,
                        min_n=n,
                        min_k=k,
                    )
                    for g in [
                        [text_g, data_manager, 1, "arxiv"],
                        [g_cora, cora_fsdm, 0, "cora"],
                    ]
                    for n in [5, 3]
                    for k in k_shot_lst
                ]
                test_data = [
                    make_data(
                        g[0],
                        g[1],
                        g[2],
                        f"test_fs{n}{k}_{g[3]}",
                        True,
                        class_emb_flag=class_emb_flag,
                        min_n=n,
                        min_k=k,
                    )
                    for g in [
                        [text_g, data_manager, 2, "arxiv"],
                        [g_cora, cora_fsdm, 0, "cora"],
                    ]
                    for n in [5, 3]
                    for k in k_shot_lst
                ]

                split_data = {
                    "train": train_data,
                    "val": valid_data,
                    "test": test_data,
                }
                constructed_data.append(split_data)

        if name.startswith("fskg_"):
            (
                [n_way, min_n],
                [k_shot, min_k],
                q_query,
                fs_task_num,
                class_emb_flag,
            ) = fs_setting[name]

            class_split_lst = None
            class_split_ratio = None

            dataset = FB15K237OFADataset("FB15K237", sentence_encoder=encoder)
            train_g = dataset.data
            train_g.x = train_g.x_text_feat
            train_g.prompt_edge_feat = dataset.prompt_edge_feat
            converted_triplet = dataset.get_idx_split()
            edges = torch.cat(
                [
                    torch.tensor(converted_triplet["train"][0]).T,
                    torch.tensor(converted_triplet["valid"][0]).T,
                    torch.tensor(converted_triplet["test"][0]).T,
                ],
                dim=-1,
            )
            train_g.edge_index = edges
            edge_labels = torch.cat(
                [
                    torch.tensor(converted_triplet["train"][1]),
                    torch.tensor(converted_triplet["valid"][1]),
                    torch.tensor(converted_triplet["test"][1]),
                ]
            )
            train_g.y = edge_labels
            train_adj = None
            train_datamanager = FewShotDataManager(
                train_g, n_way, k_shot, q_query
            )
            TRAIN_LABELS = [
                44,
                109,
                9,
                144,
                106,
                73,
                172,
                52,
                211,
                113,
                226,
                191,
                29,
                148,
                111,
                98,
                21,
                186,
                228,
                119,
                185,
                66,
                199,
                112,
                174,
                72,
                207,
                85,
                200,
                115,
                84,
                60,
                125,
                61,
                12,
                99,
                149,
                27,
                78,
                181,
                10,
                128,
                30,
                42,
                45,
                100,
                89,
                4,
                180,
                161,
                230,
                5,
                130,
                192,
                3,
                15,
                175,
                131,
                156,
                145,
                190,
                110,
                219,
                2,
                222,
                203,
                79,
                122,
                81,
                47,
                206,
                164,
                176,
                38,
                13,
                77,
                223,
                64,
                123,
                0,
                146,
                105,
                193,
                37,
                157,
                153,
                82,
                34,
                124,
                187,
                197,
                103,
                201,
                232,
                32,
                215,
                165,
                104,
                48,
                171,
                177,
                76,
                120,
                196,
                132,
                209,
                135,
                19,
                212,
                213,
                208,
                169,
                133,
                65,
                170,
                235,
                46,
                51,
                121,
                204,
                134,
                167,
                126,
                33,
                138,
                162,
                95,
                63,
                229,
                184,
                36,
                182,
                101,
                127,
                152,
                140,
                54,
                16,
                155,
                102,
                80,
                220,
            ]

            VAL_LABELS = [
                168,
                14,
                83,
                205,
                142,
                69,
                236,
                118,
                159,
                231,
                136,
                43,
                18,
                68,
                53,
                90,
                94,
                41,
                93,
                116,
                195,
                225,
                25,
                216,
                198,
                74,
                58,
                210,
                17,
                49,
                147,
                92,
                158,
                160,
                75,
                141,
                20,
                96,
                31,
                137,
                117,
                11,
                67,
                214,
                88,
                91,
                24,
            ]

            TEST_LABELS = [
                97,
                218,
                227,
                86,
                217,
                39,
                202,
                87,
                221,
                178,
                40,
                194,
                1,
                71,
                150,
                114,
                56,
                107,
                224,
                179,
                166,
                183,
                50,
                143,
                234,
                154,
                129,
                59,
                55,
                23,
                7,
                8,
                108,
                151,
                22,
                139,
                233,
                173,
                26,
                188,
                35,
                57,
                62,
                70,
                189,
                6,
                28,
                163,
            ]
            class_split_lst = [
                TRAIN_LABELS,
                VAL_LABELS,
                TEST_LABELS,
                TRAIN_LABELS + VAL_LABELS + TEST_LABELS,
            ]
            fs_edges = []
            fs_edge_types = []
            for classes in class_split_lst:
                fs_mask = torch.tensor(
                    [item in classes for item in edge_labels]
                )
                fs_edges.append(edges[:, fs_mask])
                fs_edge_types.append(edge_labels[fs_mask])

            if "wn" in name:
                dataset = WN18RROFADataset("WN18RR", sentence_encoder=encoder)
                class_split_lst = [
                    [
                        0,
                        2,
                        4,
                        5,
                    ],
                    [
                        1,
                        8,
                        9,
                    ],
                    [
                        3,
                        6,
                        7,
                    ],
                ]
                class_split_ratio = [4, 3, 3]
            else:
                raise ValueError("No dataset for test.")

            text_g = dataset.data
            text_g.x = text_g.x_text_feat
            text_g.prompt_edge_feat = dataset.prompt_edge_feat
            converted_triplet = dataset.get_idx_split()
            text_edges = torch.cat(
                [
                    torch.tensor(converted_triplet["train"][0]).T,
                    torch.tensor(converted_triplet["valid"][0]).T,
                    torch.tensor(converted_triplet["test"][0]).T,
                ],
                dim=-1,
            )
            text_g.edge_index = text_edges
            text_edge_labels = torch.cat(
                [
                    torch.tensor(converted_triplet["train"][1]),
                    torch.tensor(converted_triplet["valid"][1]),
                    torch.tensor(converted_triplet["test"][1]),
                ]
            )
            text_g.y = text_edge_labels
            data_manager = FewShotDataManager(
                text_g,
                n_way,
                k_shot,
                q_query,
                class_split_ratio=class_split_ratio,
                class_split_lst=class_split_lst,
            )

            def make_data(
                g,
                data_manager,
                split_name: int,
                state_name,
                undirected_flag,
                edges,
                fs_edges,
                fs_edge_types,
                class_emb_flag=True,
                train_flag=False,
                adj=None,
                random_flag=False,
                min_n=None,
                min_k=None,
            ):
                if class_emb_flag:
                    class_emb = g.edge_label_feat
                else:
                    class_emb = g.prompt_text_feat.repeat(
                        len(g.edge_label_feat), 1
                    )
                if k_shot > 0:
                    data_class = FewShotKGDataset
                elif k_shot == 0:
                    data_class = ZeroShotKGDataset
                else:
                    raise ValueError("check the k_shot value")
                dataset = data_class(
                    pyg_graph=g,
                    class_emb=class_emb,
                    data_idx=torch.zeros(50),
                    n_way=n_way,
                    k_shot=k_shot,
                    q_query=q_query,
                    datamanager=data_manager,
                    mode=split_name,
                    edges=edges,
                    fs_edges=fs_edges,
                    fs_edge_types=fs_edge_types,
                    hop=2,
                    prompt_feat=g.prompt_text_feat,
                    to_undirected=undirected_flag,
                    single_prompt_edge=True,
                    adj=adj,
                    random_flag=random_flag,
                    min_n=min_n,
                    min_k=min_k,
                )
                if train_flag:
                    return dataset
                return DataWithMeta(
                    dataset,
                    batch_size=fs_task_num,
                    sample_size=sample_size,
                    metric="acc",
                    state_name=state_name,
                    classes=min_n,
                    meta_data={"eval_func": classification_single_func},
                )

            if k_shot > 0:
                k_shot_lst = [1, 3, 5]
            else:
                k_shot_lst = [0]

            train_data = [
                make_data(
                    train_g,
                    train_datamanager,
                    0,
                    f"train_fs_fb",
                    True,
                    train_g.edge_index,
                    fs_edges[0],
                    fs_edge_types[0],
                    class_emb_flag=class_emb_flag,
                    train_flag=True,
                    random_flag=True,
                    min_n=min_n,
                    min_k=min_k,
                )
            ]
            valid_data = [
                make_data(
                    g[0],
                    g[1],
                    g[2],
                    f"val_fs{n}{k}_{g[3]}",
                    True,
                    g[4],
                    g[5],
                    g[6],
                    class_emb_flag=class_emb_flag,
                    min_n=n,
                    min_k=k,
                )
                for g in [
                    [
                        train_g,
                        train_datamanager,
                        1,
                        "fb",
                        train_g.edge_index,
                        fs_edges[1],
                        fs_edge_types[1],
                    ],
                    [
                        text_g,
                        data_manager,
                        3,
                        "wn",
                        text_g.edge_index,
                        text_g.edge_index,
                        text_g.y,
                    ],
                ]
                for n in [5, 3, 10]
                for k in k_shot_lst
            ]
            test_data = [
                make_data(
                    g[0],
                    g[1],
                    g[2],
                    f"test_fs{n}{k}_{g[3]}",
                    True,
                    g[4],
                    g[5],
                    g[6],
                    class_emb_flag=class_emb_flag,
                    min_n=n,
                    min_k=k,
                )
                for g in [
                    [
                        train_g,
                        train_datamanager,
                        2,
                        "fb",
                        train_g.edge_index,
                        fs_edges[2],
                        fs_edge_types[2],
                    ],
                    [
                        text_g,
                        data_manager,
                        3,
                        "wn",
                        text_g.edge_index,
                        text_g.edge_index,
                        text_g.y,
                    ],
                ]
                for n in [5, 3, 10]
                for k in k_shot_lst
            ]

            split_data = {
                "train": train_data,
                "val": valid_data,
                "test": test_data,
            }
            constructed_data.append(split_data)

    states = ["train", "test", "val"]
    collate_dataset = {}
    for s in states:
        dataset_list = []
        for ds in constructed_data:
            dataset_list += ds[s]
        collate_dataset[s] = dataset_list
    return collate_dataset


def main(params):
    encoder = SentenceEncoder("ST")

    # fs_setting = {
    #     "fs_arxiv_5_3": [5, 3, 1, 1, False],
    #     # "fs_arxiv_5_1": [5, 1, 1, 1, False],
    #     # "fs_arxiv_5_5": [5, 5, 1, 1, False],
    #     # "fs_arxiv_3_3": [3, 3, 1, 1, False],
    #     # "fs_arxiv_3_1": [3, 1, 1, 1, False],
    #     # "fs_arxiv_3_5": [3, 5, 1, 1, False],
    #     # "fs_arxiv_3_0": [3, 0, 1, 1, True],
    #     "fs_arxiv_5_0": [5, 0, 1, 1, True],
    #     # "fskg_fb_wn_5_3": [5, 3, 1, 1, False],
    #     # "fskg_fb_wn_5_1": [5, 1, 1, 1, False],
    #     # "fskg_fb_wn_5_5": [5, 5, 1, 1, False],
    #     # "fskg_fb_wn_10_3": [10, 3, 1, 1, False],
    #     # "fskg_fb_wn_10_1": [10, 1, 1, 1, False],
    #     # "fskg_fb_wn_10_5": [10, 5, 1, 1, False],
    #     # "fskg_fb_wn_5_0": [5, 0, 1, 1, True],
    #     # "fskg_fb_wn_10_0": [10, 0, 1, 1, True],
    # }
    fs_setting = {
        "fs_arxiv_1": [[5, 3], [5, 1], 1, 1, False],
        "fs_arxiv_0": [[5, 3], [0, 0], 1, 1, True],
        "fskg_fb_wn_1": [[10, 3], [5, 1], 1, 1, False],
        "fskg_fb_wn_0": [[10, 3], [0, 0], 1, 1, True],
    }
    dataset_list = list(fs_setting.keys())
    fs_ds_multiple = [50, 50, 20, 20]
    min_ratio = [50, 0, 50, 0, 20.0, 20.0]

    e2e_data_list = [
        "molzero",
        "molfew",
    ]
    mol_ds_multiple = [0.05] * 2
    mol_ds_min_ratio = [0.05] * 2

    collate_dataset = data_construct(
        dataset_list + e2e_data_list,
        encoder,
        batch_size=params.eval_batch_size,
        sample_size=params.eval_sample_size,
        walk_length=params.rwpe,
        fs_setting=fs_setting,
    )

    out_dim = 768 + (params.rwpe if params.rwpe is not None else 0)
    # out_dim = 768

    def make_data(data, b_size, sample_size):
        return DataWithMeta(data, b_size, sample_size=sample_size)

    if hasattr(params, "d_multiple"):
        data_multiple = [float(a) for a in params.d_multiple.split(",")]
    else:
        data_multiple = [2, 2, 0.3, 2, 0.5, 0.4, 0.3, 2, 1, 2, 3]

    if hasattr(params, "d_min_ratio"):
        min_ratio = [float(a) for a in params.d_min_ratio.split(",")]
    else:
        min_ratio = [0.5, 0.5, 0.05, 1, 0.1, 0.1, 0.03, 1, 0.2, 0.2, 1]

    train_data = MultiDataset(
        collate_dataset["train"],
        dataset_multiple=fs_ds_multiple + mol_ds_multiple,
        # dataset_multiple=1,
        patience=3,
        window_size=5,
        min_ratio=min_ratio + mol_ds_min_ratio,
        # min_ratio=1.0,
    )

    text_dataset = {
        "train": make_data(
            train_data,
            params.batch_size,
            params.train_sample_size,
        ),
        "val": collate_dataset["val"],
        "test": collate_dataset["test"],
    }
    params.datamodule = DataModule(
        text_dataset, num_workers=params.num_workers
    )

    eval_data = text_dataset["val"] + text_dataset["test"]
    val_state = [dt.state_name for dt in text_dataset["val"]]
    test_state = [dt.state_name for dt in text_dataset["test"]]
    eval_state = val_state + test_state
    eval_metric = [dt.metric for dt in eval_data]
    eval_funcs = [dt.meta_data["eval_func"] for dt in eval_data]
    loss = torch.nn.BCEWithLogitsLoss()
    evlter = []
    for dt in eval_data:
        if dt.metric == "acc":
            evlter.append(Accuracy(task="multiclass", num_classes=dt.classes))
        elif dt.metric == "auc":
            evlter.append(AUROC(task="binary"))
        elif dt.metric == "apr":
            evlter.append(MultiApr(num_labels=dt.classes))
        elif dt.metric == "aucmulti":
            evlter.append(MultiAuc(num_labels=dt.classes))
    metrics = EvalKit(
        eval_metric,
        evlter,
        loss,
        eval_funcs,
        flat_binary_func,
        eval_mode="max",
        exp_prefix="",
        eval_state=eval_state,
        val_monitor_state=val_state[0],
        test_monitor_state=test_state[0],
    )
    # gnn = PyGGIN(params.num_layers, 768, 768)
    # gnn = PyGRGCN(params.num_layers, 3, 768, 768)
    # gnn = PyGGINE(params.num_layers, 768, 768, 768)
    gnn = PyGRGCNEdge(
        params.num_layers,
        5,
        out_dim,
        out_dim,
        drop_ratio=params.dropout,
        JK=params.JK,
    )
    bin_model = BinGraphAttModel if params.JK == "none" else BinGraphModel
    model = bin_model(gnn, out_dim, 1, add_rwpe=params.rwpe)
    optimizer = torch.optim.Adam(
        model.parameters(), lr=params.lr, weight_decay=params.l2
    )
    lr_scheduler = {
        "scheduler": torch.optim.lr_scheduler.StepLR(optimizer, 15, 0.5),
        "interval": "epoch",
        "frequency": 1,
    }

    exp_config = ExpConfig(
        "",
        optimizer,
        dataset_callback=train_data.update,
        lr_scheduler=lr_scheduler,
    )
    exp_config.val_state_name = val_state
    exp_config.test_state_name = test_state

    pred_model = GraphPredLightning(exp_config, model, metrics)

    wandb_logger = WandbLogger(
        project=params.log_project,
        name=params.exp_name,
        save_dir=params.exp_dir,
        offline=params.offline_log,
    )

    val_res, test_res = lightning_fit(
        wandb_logger,
        pred_model,
        params.datamodule,
        metrics,
        params.num_epochs,
        save_model=True,
        load_best=False,
        reload_freq=1,
        # profiler="simple",
        # accelerator="cpu",
    )


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="rl")

    parser.add_argument(
        "opts",
        default=[],
        nargs=argparse.REMAINDER,
        help="Modify config options using the command-line",
    )

    parser.add_argument("--n_way", type=int, default=5)
    parser.add_argument("--k_shot", type=int, default=3)
    parser.add_argument("--q_query", type=int, default=3)
    parser.add_argument(
        "--fs_task_num",
        type=int,
        default=5,
        help="Number of tasks for few-shot training.",
    )

    params = parser.parse_args()
    configs = []
    configs.append(
        load_yaml(
            os.path.join(
                os.path.dirname(__file__), "configs", "default_config.yaml"
            )
        )
    )
    print(configs)
    # Add for few-shot parameters
    configs.append(params.__dict__)

    mod_params = combine_dict(*configs)
    mod_params = merge_mod(mod_params, params.opts)
    setup_exp(mod_params)

    params = SimpleNamespace(**mod_params)
    set_random_seed(params.seed)

    torch.set_float32_matmul_precision("high")
    params.log_project = "fs_cdm"
    main(params)

run_fs.py
LechengKong/OneForAll
import torch
import torch_geometric as pyg

from data.arxiv.gen_data import ArxivOFADataset
from data.Cora.gen_data import CoraOFADataset
from data.Pubmed.gen_data import PubmedOFADataset
from data.WN18RR.gen_data import WN18RROFADataset
from data.FB15K237.gen_data import FB15K237OFADataset
from data.wikics.gen_data import WikiCSOFADataset
from data.chemblpre.gen_data import CHEMBLPREOFADataset
from data.chempcba.gen_data import CHEMPCBAOFADataset
from data.chemhiv.gen_data import CHEMHIVOFADataset

from ofa_datasets import (
    GraphListDataset,
    SubgraphDataset,
    MultiDataset,
    GraphListHierDataset,
    SubgraphHierDataset,
    GraphListHierFSDataset,
    GraphListHierFixDataset,
    SubgraphLinkHierDataset,
    SubgraphKGHierDataset,
    FewShotNCDataset,
    FewShotKGDataset,
    ZeroShotNCDataset,
    ZeroShotKGDataset,
    SubgraphNopromptDataset,
    GraphListNopromptDataset,
)

from gp.utils.utils import k_fold_ind, k_fold2_split
from gp.lightning.data_template import DataWithMeta

from gp.lightning.metric import (
    binary_auc_func,
    flat_binary_func,
    classification_func,
    EvalKit,
)
from utils import (
    binary_apr_func,
    binary_auc_multi_func,
    binary_single_auc_func,
)

name2dataset = {
    "arxiv": ArxivOFADataset,
    "cora": CoraOFADataset,
    "pubmed": PubmedOFADataset,
    "WN18RR": WN18RROFADataset,
    "FB15K237": FB15K237OFADataset,
    "wikics": WikiCSOFADataset,
    "chemblpre": CHEMBLPREOFADataset,
    "chempcba": CHEMPCBAOFADataset,
    "chemhiv": CHEMHIVOFADataset,
}


def ArxivSplitter(dataset):
    text_g = dataset.data
    kfold = k_fold_ind(text_g.y, 10)
    text_split = k_fold2_split(kfold, len(text_g.y))[0]
    split = {}
    split["train"] = text_split[0]
    split["valid"] = text_split[1]
    split["test"] = text_split[2]
    return split


def CiteSplitter(dataset):
    text_g = dataset.data
    split = {
        "train": text_g.train_masks[0].nonzero(as_tuple=True)[0],
        "valid": text_g.val_masks[0].nonzero(as_tuple=True)[0],
        "test": text_g.test_masks[0].nonzero(as_tuple=True)[0],
    }
    return split


def CiteLinkSplitter(dataset):
    text_g = dataset.data
    text_g.x = text_g.x_text_feat
    text_g.prompt_edge_feat = dataset.prompt_edge_feat
    edges = text_g.edge_index
    edge_perm = torch.randperm(len(edges[0]))
    train_offset = int(len(edge_perm) * 0.85)
    val_offset = int(len(edge_perm) * 0.9)
    edge_indices = {
        "train": edge_perm[:train_offset],
        "valid": edge_perm[train_offset:val_offset],
        "test": edge_perm[val_offset:],
    }
    return edge_indices


def KGSplitter(dataset):
    converted_triplet = dataset.get_idx_split()
    return converted_triplet


def WikiSplitter(dataset):
    text_g = dataset.data
    wiki_split_idx = 0
    split = {
        "train": torch.where(text_g.train_mask[:, wiki_split_idx])[0].numpy(),
        "valid": torch.where(text_g.val_mask[:, wiki_split_idx])[0].numpy(),
        "test": torch.where(text_g.test_mask)[0].numpy(),
    }
    return split


def MolSplitter(dataset):
    return dataset.get_idx_split()


name2splitter = {
    "arxiv": ArxivSplitter,
    "cora_node": CiteSplitter,
    "pubmed_node": CiteSplitter,
    "cora_link": CiteLinkSplitter,
    "pubmed_link": CiteLinkSplitter,
    "WN18RR": KGSplitter,
    "FB15K237": KGSplitter,
    "wikics": WikiSplitter,
    "chemblpre": MolSplitter,
    "chempcba": MolSplitter,
    "chemhiv": MolSplitter,
}


def LinkConstructGraph(dataset, split):
    text_g = dataset.data
    edges = text_g.edge_index
    graph_dict = text_g.to_dict()
    graph_dict["edge_index"] = edges[:, split["train"]]
    train_graph = pyg.data.Data(**graph_dict)
    return train_graph


def make_data(
    name, data, split_name, metric, eval_func, num_classes, **kwargs
):
    return DataWithMeta(
        data,
        kwargs["batch_size"],
        sample_size=kwargs["sample_size"],
        metric=metric,
        state_name=split_name + "_" + name,
        classes=num_classes,
        meta_data={"eval_func": eval_func},
    )


def ConstructNodeCls(
    name, dataset, split, split_name, to_bin_cls_func, **kwargs
):
    text_g = dataset.data

    return SubgraphHierDataset(
        text_g,
        text_g.label_text_feat,
        split[split_name],
        prompt_feat=text_g.prompt_text_feat,
        to_undirected=True,
        trim_class_func=to_bin_cls_func,
        walk_length=kwargs["walk_length"],
    )


def ConstructLinkCls(
    name, dataset, split, split_name, to_bin_cls_func, **kwargs
):
    text_g = dataset.data
    edges = text_g.edge_index
    train_graph = kwargs["global_data"]

    return SubgraphLinkHierDataset(
        train_graph,
        train_graph.edge_label_feat,
        edges.T[split[split_name]].numpy(),
        prompt_feat=train_graph.prompt_text_edge_feat,
        to_undirected=True,
        hop=3,
        remove_edge=kwargs["remove_edge"],
        trim_class_func=to_bin_cls_func,
        walk_length=kwargs["walk_length"],
    )


def ConstructKG(name, dataset, split, split_name, to_bin_cls_func, **kwargs):
    text_g = dataset.data

    return SubgraphKGHierDataset(
        text_g,
        text_g.edge_label_feat,
        split[split_name],
        prompt_feat=text_g.prompt_text_feat,
        to_undirected=True,
        hop=2,
        remove_edge=kwargs["remove_edge"],
        trim_class_func=to_bin_cls_func,
        walk_length=kwargs["walk_length"],
    )


def ConstructMolCls(
    name, dataset, split, split_name, to_bin_cls_func, **kwargs
):

    return GraphListHierDataset(
        dataset,
        dataset.label_text_feat,
        dataset.prompt_edge_feat,
        dataset.prompt_text_feat,
        split[split_name],
        trim_class_func=to_bin_cls_func,
        single_prompt_edge=True,
        walk_length=kwargs["walk_length"],
    )


def ConstructMolFSTrain(
    name, dataset, split, split_name, to_bin_cls_func, **kwargs
):
    classes = dataset.y.view(len(dataset), -1)[split["train"]]

    return GraphListHierFSDataset(
        dataset,
        dataset.label_text_feat,
        dataset.prompt_edge_feat,
        dataset.prompt_text_feat,
        split[split_name],
        trim_class_func=to_bin_cls_func,
        single_prompt_edge=True,
        walk_length=kwargs["walk_length"],
        class_ind=classes,
    )


def process_pth_label(embs, label):
    binary_rep = torch.zeros((1, len(embs)))
    binary_rep[0, label] = 1
    return label.view(1, -1), embs, binary_rep


def process_multi_label(embs, label):
    valid_idx = label == label
    # valid_idx = torch.zeros_like(classes, dtype=torch.bool)
    return (
        torch.tensor([[0]]),
        embs[valid_idx.view(-1)].detach().clone(),
        label[:, valid_idx.view(-1)].detach().clone(),
    )


def process_int_label(embs, label):
    binary_rep = torch.zeros((1, len(embs)))
    binary_rep[0, label] = 1
    return torch.tensor([label]).view(1, -1), embs, binary_rep


none_process_label = None


task_config_lookup = {
    "arxiv": {
        "dataset_name": "arxiv",
        "dataset_splitter": "ArxivSplitter",
        "preprocess": None,
        "construct": "ConstructNodeCls",
        "args": {"walk_length": None},
        "process_label_func": "process_pth_label",
        "eval_set_constructs": [
            {
                "stage": "valid",
                "split_name": "valid",
            },
            {
                "stage": "test",
                "split_name": "test",
            },
            {
                "stage": "test",
                "split_name": "train",
            },
        ],
        "eval_metric": "acc",
        "eval_func": "classification_func",
        "num_classes": 40,
    },
    "cora_link": {
        "dataset_name": "cora",
        "dataset_splitter": "CiteLinkSplitter",
        "preprocess": "LinkConstructGraph",
        "construct": "ConstructLinkCls",
        "args": {"remove_edge": True, "walk_length": None},
        "process_label_func": "process_int_label",
        "eval_set_constructs": [
            {
                "stage": "valid",
                "split_name": "valid",
                "args": {"remove_edge": False, "walk_length": None},
            },
            {
                "stage": "test",
                "split_name": "test",
                "args": {"remove_edge": False, "walk_length": None},
            },
            {
                "stage": "test",
                "split_name": "train",
                "args": {"remove_edge": True, "walk_length": None},
            },
        ],
        "eval_metric": "auc",
        "eval_func": "binary_auc_func",
        "num_classes": 2,
    },
    "cora_node": {
        "dataset_name": "cora",
        "dataset_splitter": "CiteSplitter",
        "preprocess": None,
        "construct": "ConstructNodeCls",
        "args": {"walk_length": None},
        "process_label_func": "process_int_label",
        "eval_set_constructs": [
            {
                "stage": "valid",
                "split_name": "valid",
            },
            {
                "stage": "test",
                "split_name": "test",
            },
            {
                "stage": "test",
                "split_name": "train",
            },
        ],
        "eval_metric": "acc",
        "eval_func": "classification_func",
        "num_classes": 7,
    },
    "pubmed_link": {
        "dataset_name": "pubmed",
        "dataset_splitter": "CiteLinkSplitter",
        "preprocess": "LinkConstructGraph",
        "construct": "ConstructLinkCls",
        "args": {"remove_edge": True, "walk_length": None},
        "process_label_func": "process_int_label",
        "eval_set_constructs": [
            {
                "stage": "valid",
                "split_name": "valid",
                "args": {"remove_edge": False, "walk_length": None},
            },
            {
                "stage": "test",
                "split_name": "test",
                "args": {"remove_edge": False, "walk_length": None},
            },
            {
                "stage": "test",
                "split_name": "train",
                "args": {"remove_edge": True, "walk_length": None},
            },
        ],
        "eval_metric": "auc",
        "eval_func": "binary_auc_func",
        "num_classes": 2,
    },
    "pubmed_node": {
        "dataset_name": "pubmed",
        "dataset_splitter": "CiteSplitter",
        "preprocess": None,
        "construct": "ConstructNodeCls",
        "args": {"walk_length": None},
        "process_label_func": "process_int_label",
        "eval_set_constructs": [
            {
                "stage": "valid",
                "split_name": "valid",
            },
            {
                "stage": "test",
                "split_name": "test",
            },
            {
                "stage": "test",
                "split_name": "train",
            },
        ],
        "eval_metric": "acc",
        "eval_func": "classification_func",
        "num_classes": 3,
    },
    "WN18RR": {
        "dataset_name": "WN18RR",
        "dataset_splitter": "KGSplitter",
        "preprocess": None,
        "construct": "ConstructKG",
        "args": {"remove_edge": True, "walk_length": None},
        "process_label_func": "process_int_label",
        "eval_set_constructs": [
            {
                "stage": "valid",
                "split_name": "valid",
                "args": {"remove_edge": False, "walk_length": None},
            },
            {
                "stage": "test",
                "split_name": "test",
                "args": {"remove_edge": False, "walk_length": None},
            },
            {
                "stage": "test",
                "split_name": "train",
                "args": {"remove_edge": True, "walk_length": None},
            },
        ],
        "eval_metric": "acc",
        "eval_func": "classification_func",
        "num_classes": 11,
    },
    "FB15K237": {
        "dataset_name": "FB15K237",
        "dataset_splitter": "KGSplitter",
        "preprocess": None,
        "construct": "ConstructKG",
        "args": {"remove_edge": True, "walk_length": None},
        "process_label_func": "process_int_label",
        "eval_set_constructs": [
            {
                "stage": "valid",
                "split_name": "valid",
                "args": {"remove_edge": False, "walk_length": None},
            },
            {
                "stage": "test",
                "split_name": "test",
                "args": {"remove_edge": False, "walk_length": None},
            },
            {
                "stage": "test",
                "split_name": "train",
                "args": {"remove_edge": True, "walk_length": None},
            },
        ],
        "eval_metric": "acc",
        "eval_func": "classification_func",
        "num_classes": 237,
    },
    "wikics": {
        "dataset_name": "wikics",
        "dataset_splitter": "WikiSplitter",
        "preprocess": None,
        "construct": "ConstructNodeCls",
        "args": {"walk_length": None},
        "process_label_func": "process_pth_label",
        "eval_set_constructs": [
            {
                "stage": "valid",
                "split_name": "valid",
            },
            {
                "stage": "test",
                "split_name": "test",
            },
            {
                "stage": "test",
                "split_name": "train",
            },
        ],
        "eval_metric": "acc",
        "eval_func": "classification_func",
        "num_classes": 10,
    },
    "chemblpre": {
        "dataset_name": "chemblpre",
        "dataset_splitter": "MolSplitter",
        "preprocess": None,
        "construct": "ConstructMolCls",
        "args": {"walk_length": None},
        "process_label_func": "process_multi_label",
        "eval_set_constructs": [
            {
                "stage": "valid",
                "split_name": "valid",
                "process_label_func": "none_process_label",
            },
            {
                "stage": "test",
                "split_name": "test",
                "process_label_func": "none_process_label",
            },
            {
                "stage": "test",
                "split_name": "train",
                "process_label_func": "none_process_label",
            },
        ],
        "eval_metric": "apr",
        "eval_func": "binary_apr_func",
        "num_classes": 1296,
    },
    "chempcba": {
        "dataset_name": "chempcba",
        "dataset_splitter": "MolSplitter",
        "preprocess": None,
        "construct": "ConstructMolCls",
        "args": {"walk_length": None},
        "process_label_func": "process_multi_label",
        "eval_set_constructs": [
            {
                "stage": "valid",
                "split_name": "valid",
                "process_label_func": "none_process_label",
            },
            {
                "stage": "test",
                "split_name": "test",
                "process_label_func": "none_process_label",
            },
            {
                "stage": "test",
                "split_name": "train",
                "process_label_func": "none_process_label",
            },
        ],
        "eval_metric": "apr",
        "eval_func": "binary_apr_func",
        "num_classes": 128,
    },
    "chemhiv": {
        "dataset_name": "chemhiv",
        "dataset_splitter": "MolSplitter",
        "preprocess": None,
        "construct": "ConstructMolCls",
        "args": {"walk_length": None},
        "process_label_func": "process_pth_label",
        "eval_set_constructs": [
            {
                "stage": "valid",
                "split_name": "valid",
                "process_label_func": "none_process_label",
            },
            {
                "stage": "test",
                "split_name": "test",
                "process_label_func": "none_process_label",
            },
            {
                "stage": "test",
                "split_name": "train",
                "process_label_func": "none_process_label",
            },
        ],
        "eval_metric": "auc",
        "eval_func": "binary_auc_func",
        "num_classes": 2,
    },
}


class TaskConstructor:
    def __init__(self, tasks, encoder, batch_size=256, sample_size=-1):
        self.tasks = tasks
        self.batch_size = batch_size
        self.sample_size = sample_size
        self.dataset = {}
        self.train_set = []
        self.valid_dm_set = []
        self.test_dm_set = []

        for task in self.tasks:
            config = task_config_lookup[task]
            data = config["dataset_name"]
            if data not in self.dataset and data in name2dataset:
                self.dataset[data] = name2dataset[data](
                    data, sentence_encoder=encoder
                )

            split = globals()[config["dataset_splitter"]](self.dataset[data])
            if config["preprocess"] is not None:
                global_data = globals()[config["preprocess"]](
                    self.dataset[data], split
                )
            else:
                global_data = None

            train_data = globals()[config["construct"]](
                task,
                self.dataset[data],
                split,
                "train",
                globals()[config["process_label_func"]],
                global_data=global_data,
                **config["args"],
            )
            self.train_set.append(train_data)

            for eval_construct_config in config["eval_set_constructs"]:
                if "process_label_func" in eval_construct_config:
                    trim_class_func = globals()[
                        eval_construct_config["process_label_func"]
                    ]
                else:
                    trim_class_func = globals()[config["process_label_func"]]

                if "args" in eval_construct_config:
                    eval_args = eval_construct_config["args"]
                else:
                    eval_args = config["args"]

                if "construct" in eval_construct_config:
                    construct = globals()[eval_construct_config["construct"]]
                else:
                    construct = globals()[config["construct"]]
                eval_data = construct(
                    task,
                    self.dataset[data],
                    split,
                    eval_construct_config["split_name"],
                    trim_class_func,
                    global_data=global_data,
                    **eval_args,
                )

                dm_data = make_data(
                    data,
                    eval_data,
                    eval_construct_config["split_name"],
                    config["eval_metric"],
                    globals()[config["eval_func"]],
                    config["num_classes"],
                    batch_size=self.batch_size,
                    sample_size=self.sample_size,
                )

                if eval_construct_config["stage"] == "valid":
                    self.valid_dm_set.append(dm_data)
                else:
                    self.test_dm_set.append(dm_data)

    def make_train_data(self, multiple, min_ratio):
        train_data = MultiDataset(
            self.train_set,
            dataset_multiple=multiple,
            patience=3,
            window_size=5,
            min_ratio=min_ratio,
        )
        return train_data

    def make_full_dm_list(self, multiple, min_ratio, train_data=None):
        text_dataset = {
            "train": DataWithMeta(
                self.make_train_data(multiple, min_ratio)
                if not train_data
                else train_data,
                self.batch_size,
                sample_size=self.sample_size,
            ),
            "val": self.valid_dm_set,
            "test": self.test_dm_set,
        }
        return text_dataset

task_constructor.py
LechengKong/OneForAll
import os
import torch
from sentence_transformers import SentenceTransformer
from transformers import (
    LlamaForCausalLM,
    LlamaTokenizer,
    LlamaModel,
    LlamaConfig,
)
from torchmetrics import AveragePrecision, AUROC
import numpy as np
from torch_geometric.utils import (
    to_scipy_sparse_matrix,
    scatter,
)


class SentenceEncoder:
    def __init__(
        self, name, root="cache_data/model", batch_size=512, multi_gpu=False
    ):
        self.batch_size = batch_size
        self.multi_gpu = multi_gpu
        if name == "ST":
            self.model = SentenceTransformer(
                "multi-qa-distilbert-cos-v1",
                device="cuda:0",
                cache_folder=root,
            )
        elif name == "llama2":
            model_path = os.path.join(os.getcwd(), "llama-2-7b")
            self.tokenizer = LlamaTokenizer.from_pretrained(
                model_path, device="cuda:0"
            )
            self.model = LlamaModel.from_pretrained(model_path).to("cuda:0")
        elif name == "roberta":
            self.model = SentenceTransformer(
                "sentence-transformers/roberta-base-nli-stsb-mean-tokens",
                device="cuda:0",
                cache_folder=root,
            )
        else:
            raise ValueError(f"Unknown language model: {name}.")

    def encode(self, texts, to_tensor=True):
        if self.multi_gpu:
            # Start the multi-process pool on all available CUDA devices
            pool = self.model.start_multi_process_pool()
            embeddings = self.model.encode_multi_process(
                texts,
                pool=pool,
                batch_size=self.batch_size,
            )
            embeddings = torch.from_numpy(embeddings)
        else:
            embeddings = self.model.encode(
                texts,
                batch_size=self.batch_size,
                show_progress_bar=True,
                convert_to_tensor=to_tensor,
                convert_to_numpy=not to_tensor,
            )
        return embeddings


def binary_single_auc_func(func, output, batch):
    output = output.view(-1, batch.num_classes[0])
    score = torch.sigmoid(output)
    # if len(score.unique()) == 1:
    # print(output[:20])
    label = batch.bin_labels[batch.true_nodes_mask]
    # print(score)
    # print(label)
    return func.update(score, label.view(-1, batch.num_classes[0]))


def binary_apr_func(func, output, batch):
    output = output.view(-1, batch.num_classes[0])
    score = torch.sigmoid(output)
    label = batch.bin_labels[batch.true_nodes_mask]
    return func.update(score, label.view(len(batch), -1))


def binary_auc_multi_func(func, output, batch):
    output = output.view(-1, batch.num_classes[0])
    score = torch.sigmoid(output)
    label = batch.bin_labels[batch.true_nodes_mask]
    return func.update(score, label.view(-1, batch.num_classes[0]))


def label_apr_func(func, output, batch):
    score = torch.sigmoid(output)
    return func.update(score, batch.y)


def flat_label_func(func, output, batch):
    labels = batch.y.view(-1)
    valid_ind = labels == labels
    return func(output.view(-1)[valid_ind], labels[valid_ind])


def classification_single_func(func, output, batch):
    label = batch.bin_labels[batch.true_nodes_mask].view(-1, batch.num_classes[0])
    output = output.view(-1, batch.num_classes[0])
    return func(output, torch.argmax(label, dim=-1))


class MultiApr(torch.nn.Module):
    def __init__(self, num_labels=1):
        super().__init__()
        self.metrics = torch.nn.ModuleList(
            [AveragePrecision(task="binary") for i in range(num_labels)]
        )

    def update(self, preds, targets):
        for i, met in enumerate(self.metrics):
            pred = preds[:, i]
            target = targets[:, i]
            valid_idx = target == target
            # print(pred[valid_idx])
            # print(target[valid_idx])
            met.update(pred[valid_idx], target[valid_idx].to(torch.long))

    def compute(self):
        full_val = []
        for met in self.metrics:
            try:
                res = met.compute()
                if res == res:
                    full_val.append(res)
            except BaseException:
                pass
        return torch.tensor(full_val).mean()

    def reset(self):
        for met in self.metrics:
            met.reset()


class MultiAuc(torch.nn.Module):
    def __init__(self, num_labels=1):
        super().__init__()
        self.metrics = torch.nn.ModuleList(
            [AUROC(task="binary") for i in range(num_labels)]
        )

    def update(self, preds, targets):
        for i, met in enumerate(self.metrics):
            pred = preds[:, i]
            target = targets[:, i]
            valid_idx = target == target
            # print(pred[valid_idx])
            # print(target[valid_idx])
            met.update(pred[valid_idx], target[valid_idx].to(torch.long))

    def compute(self):
        full_val = []
        for met in self.metrics:
            try:
                res = met.compute()
                if res == res:
                    full_val.append(res)
            except BaseException:
                pass
        return torch.tensor(full_val).mean()

    def reset(self):
        for met in self.metrics:
            met.reset()


def scipy_rwpe(data, walk_length):
    row, col = data.edge_index
    N = data.num_nodes

    value = data.edge_weight
    if value is None:
        value = torch.ones(data.num_edges, device=row.device)
    value = scatter(value, row, dim_size=N, reduce="sum").clamp(min=1)[row]
    value = 1.0 / value
    adj = to_scipy_sparse_matrix(
        data.edge_index, edge_attr=value, num_nodes=data.num_nodes
    )

    out = adj
    pe_list = [out.diagonal()]
    for _ in range(walk_length - 1):
        out = out @ adj
        pe_list.append(out.diagonal())
    pe = torch.tensor(np.stack(pe_list, axis=-1))

    return pe

utils.py
LechengKong/OneForAll