YuWang0103/LGGM-Text2Graph
1
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