parkererickson/LGGM-Text2Graph

1import torch2from distributions import DistributionNodes3from utils import to_dense4from torch_geometric.loader import DataLoader5from torch_geometric.data import Data6from torch_geometric.utils import remove_self_loops, to_undirected7import os8from sentence_transformers import SentenceTransformer9import random10 11 12def arrange_data(adj_matrix, cond_emb, ind):13    n_nodes = adj_matrix.shape[0]14 15    edge_index = adj_matrix.nonzero().t()16    edge_attr = torch.tensor([[0, 1] for _ in range(edge_index.shape[1])])17 18    edge_index, edge_attr = to_undirected(edge_index, edge_attr, n_nodes, reduce = 'mean')19    edge_index, edge_attr = remove_self_loops(edge_index, edge_attr)20 21    x = torch.ones((n_nodes, 1))22 23    y = torch.empty(1, 0)24    cond_emb = torch.tensor(cond_emb).unsqueeze(0)25 26    return Data(x=x, edge_index=edge_index, edge_attr=edge_attr, y=y, prompt_id = torch.tensor(ind), cond_emb = cond_emb)27 28 29 30def load_dataset_cc(dataname, batch_size, hydra_path, condition):31    domains = ['cc_high', 'cc_medium', 'cc_low']32    33    34    model = SentenceTransformer("all-MiniLM-L6-v2")35    cond_embs = model.encode(condition)36 37    38 39    train_data, val_data, test_data = [], [], []40 41    if dataname in domains: #only for test42        train_d = torch.load(f'{hydra_path}/graphs/{dataname}/train.pt')43        val_d = torch.load(f'{hydra_path}/graphs/{dataname}/val.pt')44        test_d = torch.load(f'{hydra_path}/graphs/{dataname}/test.pt')45 46        train_indices = torch.load(f'{hydra_path}/graphs/{dataname}/train_indices.pt')47        val_indices = torch.load(f'{hydra_path}/graphs/{dataname}/val_indices.pt')48        test_indices = torch.load(f'{hydra_path}/graphs/{dataname}/test_indices.pt')49 50        with open(f'{hydra_path}/graphs/{dataname}/text_prompt_order.txt', 'r') as f:51            text_prompt = f.readlines()52            text_prompt = [x.strip() for x in text_prompt]53        54        # text_prompt = ['1111111shgowhgo234o234']*1000055        print(text_prompt[0])56        text_embs = model.encode(text_prompt)57        cond_embs = torch.tensor(text_embs)58 59        train_data.extend([arrange_data(d, text_embs[ind.item()], ind.item()) for d, ind in zip(train_d, train_indices)])60        val_data.extend([arrange_data(d, text_embs[ind.item()], ind.item()) for d, ind in zip(val_d, val_indices)])61        62 63        if dataname != 'eco':64            # test_data = [arrange_data(d, text_embs[ind.item()], ind.item()) for d, ind in zip(test_d, test_indices)]65            test_data = [arrange_data(d, text_embs[ind.item()], ind.item()) for d, ind in zip(test_d, test_indices)]66        else:67            test_data = [arrange_data(d, text_embs[ind.item()], ind.item()) for d, ind in zip(train_d, train_indices)] + [arrange_data(d, text_embs[ind.item()], ind.item()) for d, ind in zip(val_data, val_indices)] + [arrange_data(d, text_embs[ind.item()], ind.item()) for d, ind in zip(test_data, test_indices)]68 69 70 71    print('Size of dataset', len(train_data), len(val_data), len(test_data))72 73    train_loader = DataLoader(train_data, batch_size = batch_size, shuffle=True)74    val_loader = DataLoader(val_data, batch_size = batch_size, shuffle=False)75    test_loader = DataLoader(test_data, batch_size=batch_size, shuffle=False)76 77    return train_loader, val_loader, test_loader, train_data, val_data, test_data, text_embs.shape[1], torch.tensor(cond_embs)78 79 80 81 82def load_dataset_deg(dataname, batch_size, hydra_path, condition):83    domains = ['deg_high', 'deg_medium', 'deg_low']84    85    86    model = SentenceTransformer("all-MiniLM-L6-v2")87    cond_embs = model.encode(condition)88 89    for domain in domains:90        if not os.path.exists(f'{hydra_path}/graphs/{domain}/train.pt'):91 92            data = torch.load(f'{hydra_path}/graphs/{domain}/{domain}.pt')93 94            #fix seed95            torch.manual_seed(0)96 97            #random permute and split98            n = len(data)99            indices = torch.randperm(n)100 101            if domain == 'eco':102                train_indices = indices[:4].repeat(50)103                val_indices = indices[4:5].repeat(50)104                test_indices = indices[5:]105            else:106                train_indices = indices[:int(0.7 * n)]107                val_indices = indices[int(0.7 * n):int(0.8 * n)]108                test_indices = indices[int(0.8 * n):]109 110            train_data = [data[_] for _ in train_indices]111            val_data = [data[_] for _ in val_indices]112            test_data = [data[_] for _ in test_indices]113 114            torch.save(train_indices, f'{hydra_path}/graphs/{domain}/train_indices.pt')115            torch.save(val_indices, f'{hydra_path}/graphs/{domain}/val_indices.pt')116            torch.save(test_indices, f'{hydra_path}/graphs/{domain}/test_indices.pt')117            118            torch.save(train_data, f'{hydra_path}/graphs/{domain}/train.pt')119            torch.save(val_data, f'{hydra_path}/graphs/{domain}/val.pt')120            torch.save(test_data, f'{hydra_path}/graphs/{domain}/test.pt')121 122 123    train_data, val_data, test_data = [], [], []124 125    if dataname in domains: #only for test126        train_d = torch.load(f'{hydra_path}/graphs/{dataname}/train.pt')127        val_d = torch.load(f'{hydra_path}/graphs/{dataname}/val.pt')128        test_d = torch.load(f'{hydra_path}/graphs/{dataname}/test.pt')129 130        train_indices = torch.load(f'{hydra_path}/graphs/{dataname}/train_indices.pt')131        val_indices = torch.load(f'{hydra_path}/graphs/{dataname}/val_indices.pt')132        test_indices = torch.load(f'{hydra_path}/graphs/{dataname}/test_indices.pt')133 134        with open(f'{hydra_path}/graphs/{dataname}/text_prompt_order.txt', 'r') as f:135            text_prompt = f.readlines()136            text_prompt = [x.strip() for x in text_prompt]137        138 139        text_embs = model.encode(text_prompt)140        cond_embs = torch.tensor(text_embs)141 142        train_data.extend([arrange_data(d, text_embs[ind.item()], ind.item()) for d, ind in zip(train_d, train_indices)])143        val_data.extend([arrange_data(d, text_embs[ind.item()], ind.item()) for d, ind in zip(val_d, val_indices)])144        145 146        if dataname != 'eco':147            test_data = [arrange_data(d, text_embs[ind.item()], ind.item()) for d, ind in zip(test_d, test_indices)]148        else:149            test_data = [arrange_data(d, text_embs[ind.item()], ind.item()) for d, ind in zip(train_d, train_indices)] + [arrange_data(d, text_embs[ind.item()], ind.item()) for d, ind in zip(val_data, val_indices)] + [arrange_data(d, text_embs[ind.item()], ind.item()) for d, ind in zip(test_data, test_indices)]150 151 152    elif dataname == 'all':153        for i, domain in enumerate(domains):154            train_d = torch.load(f'{hydra_path}/graphs/{domain}/train.pt')155            val_d = torch.load(f'{hydra_path}/graphs/{domain}/val.pt')156            test_d = torch.load(f'{hydra_path}/graphs/{domain}/test.pt')157 158            train_indices = torch.load(f'{hydra_path}/graphs/{domain}/train_indices.pt')159            val_indices = torch.load(f'{hydra_path}/graphs/{domain}/val_indices.pt')160            test_indices = torch.load(f'{hydra_path}/graphs/{domain}/test_indices.pt')161 162            # text_prompt = torch.load(f'{hydra_path}/graphs/{domain}/text_prompt_order.pt')163            164            with open(f'{hydra_path}/graphs/{domain}/text_prompt_order.txt', 'r') as f:165                text_prompt = f.readlines()166                text_prompt = [x.strip() for x in text_prompt]167 168            print(domain, text_prompt[0])169 170            text_embs = model.encode(text_prompt)171            172            train_data.extend([arrange_data(d, text_embs[ind.item()], ind.item()) for d, ind in zip(train_d, train_indices)])173            val_data.extend([arrange_data(d, text_embs[ind.item()], ind.item()) for d, ind in zip(val_d, val_indices)])174            test_data.extend([arrange_data(d, text_embs[ind.item()], ind.item())  for d, ind in zip(test_d, test_indices)])175            print(i, domain, len(train_data), len(val_data), len(test_data))176 177    print('Size of dataset', len(train_data), len(val_data), len(test_data))178 179    train_loader = DataLoader(train_data, batch_size = batch_size, shuffle=True)180    val_loader = DataLoader(val_data, batch_size = batch_size, shuffle=False)181    test_loader = DataLoader(test_data, batch_size=batch_size, shuffle=False)182 183    return train_loader, val_loader, test_loader, train_data, val_data, test_data, text_embs.shape[1], torch.tensor(cond_embs)184 185 186 187 188def init_dataset(dataname, batch_size, hydra_path, condition, transition):189    train_loader, val_loader, test_loader, train_data, val_data, test_data, cond_dims, cond_emb = load_dataset_cc(dataname, batch_size, hydra_path, condition)190 191    n_nodes = node_counts(1000, train_loader, val_loader)192    node_types = torch.tensor([1]) #No node types193    edge_types = edge_counts(train_loader)194    195    num_classes = len(node_types)196    max_n_nodes = len(n_nodes) - 1197    nodes_dist = DistributionNodes(n_nodes)198 199    print('Distribution of Number of Nodes:', n_nodes)200    print('Distribution of Node Types:', node_types)201    print('Distribution of Edge Types:', edge_types)202    203    data_loaders = {'train': train_loader, 'val': val_loader, 'test': test_loader}204 205    return data_loaders, num_classes, max_n_nodes, nodes_dist, edge_types, node_types, n_nodes, cond_dims, cond_emb206 207 208def node_counts(max_nodes_possible, train_loader, val_loader):209    #Count the distribution of graph size210    all_counts = torch.zeros(max_nodes_possible)211    212    for loader in [train_loader, val_loader]:213        for data in loader:214            unique, counts = torch.unique(data.batch, return_counts=True)215            for count in counts:216                all_counts[count] += 1217 218    max_index = max(all_counts.nonzero())219    all_counts = all_counts[:max_index + 1]220    all_counts = all_counts / all_counts.sum()221    222    return all_counts223 224def node_counts_meta(max_nodes_possible, train_data, val_data, num_classes):225    #Count the distribution of graph size226 227    all_counts = [torch.zeros(max_nodes_possible) for _ in range(num_classes)]228    229    for dataset in [train_data, val_data]:230        for data in dataset:231            all_counts[data.cond_type.item()][data.x.shape[0]] += 1232    233    for _ in range(num_classes):234        tmp = all_counts[_].nonzero()235        if len(tmp) == 0:236            max_index = 1237            all_counts[_][0] = 1238        else:239            max_index = max(tmp)240        241        all_counts[_] = all_counts[_][:max_index + 1]        242        all_counts[_] = all_counts[_] / all_counts[_].sum()243    244    return all_counts245 246 247def node_types(train_loader):248    #Count the marginal distribution of node types249    num_classes = None250    for data in train_loader:251        num_classes = data.x.shape[1]252        break253 254    counts = torch.zeros(num_classes)255 256    for i, data in enumerate(train_loader):257        counts += data.x.sum(dim=0)258 259    counts = counts / counts.sum()260    return counts261 262def edge_counts(train_loader):263    #Count the marginal distribution of edge types264    num_classes = None265    for data in train_loader:266        num_classes = data.edge_attr.shape[1]267        break268 269    d = torch.zeros(num_classes, dtype=torch.float)270 271    for i, data in enumerate(train_loader):272        unique, counts = torch.unique(data.batch, return_counts=True)273 274        all_pairs = 0275        for count in counts:276            all_pairs += count * (count - 1)277 278        279        num_edges = data.edge_index.shape[1]280        num_non_edges = all_pairs - num_edges281            282        edge_types = data.edge_attr.sum(dim=0)283        assert num_non_edges >= 0284        d[0] += num_non_edges285        d[1:] += edge_types[1:]286 287    d = d / d.sum()288    return d289 290 291def edge_counts_meta(train_data, num_classes):292    #Count the marginal distribution of edge types293    num_edge_classes = None294    for data in train_data:295        num_edge_classes = data.edge_attr.shape[1]296        break297 298    d = [torch.ones(num_edge_classes, dtype=torch.float) for _ in range(num_classes)]299 300    for i, data in enumerate(train_data):301        n_nodes = data.x.shape[0]302        303        all_pairs = n_nodes * (n_nodes - 1)304        num_edges = data.edge_index.shape[1]305        num_non_edges = all_pairs - num_edges306            307        edge_types = data.edge_attr.sum(dim=0)308        assert num_non_edges >= 0309        d[data.cond_type.item()][0] += num_non_edges310        d[data.cond_type.item()][1:] += edge_types[1:]311 312    for i, _ in enumerate(d):313        d[i] = d[i] / d[i].sum()314    315    d = torch.stack(d)316 317    return d318 319 320def compute_input_output_dims(train_loader, extra_features):321    example_batch = next(iter(train_loader))322    ex_dense, node_mask = to_dense(example_batch.x, example_batch.edge_index, example_batch.edge_attr, example_batch.batch)323 324    example_data = {'X_t': ex_dense.X, 'E_t': ex_dense.E, 'y_t': example_batch['y'], 'node_mask': node_mask}325 326    input_dims = {'X': example_batch['x'].size(1),327                  'E': example_batch['edge_attr'].size(1),328                  'y': example_batch['y'].size(1) + 1}      # + 1 due to time conditioning329    330    ex_extra_feat = extra_features(example_data)331    input_dims['X'] += ex_extra_feat.X.size(-1)332    input_dims['E'] += ex_extra_feat.E.size(-1)333    input_dims['y'] += ex_extra_feat.y.size(-1)334 335    output_dims = {'X': example_batch['x'].size(1),336                   'E': example_batch['edge_attr'].size(1),337                   'y': 0}338 339    return input_dims, output_dims340