-
Notifications
You must be signed in to change notification settings - Fork 56
/
test_ranking.py
executable file
·543 lines (400 loc) · 22.4 KB
/
test_ranking.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
import os
import random
import argparse
import logging
import json
import time
import multiprocessing as mp
import scipy.sparse as ssp
from tqdm import tqdm
import networkx as nx
import torch
import numpy as np
import dgl
def process_files(files, saved_relation2id, add_traspose_rels):
'''
files: Dictionary map of file paths to read the triplets from.
saved_relation2id: Saved relation2id (mostly passed from a trained model) which can be used to map relations to pre-defined indices and filter out the unknown ones.
'''
entity2id = {}
relation2id = saved_relation2id
triplets = {}
ent = 0
rel = 0
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
# Save the triplets corresponding to only the known relations
if triplet[1] in saved_relation2id:
data.append([entity2id[triplet[0]], entity2id[triplet[2]], saved_relation2id[triplet[1]]])
triplets[file_type] = np.array(data)
id2entity = {v: k for k, v in entity2id.items()}
id2relation = {v: k for k, v in relation2id.items()}
# Construct the list of adjacency matrix each corresponding to eeach relation. Note that this is constructed only from the train data.
adj_list = []
for i in range(len(saved_relation2id)):
idx = np.argwhere(triplets['graph'][:, 2] == i)
adj_list.append(ssp.csc_matrix((np.ones(len(idx), dtype=np.uint8), (triplets['graph'][:, 0][idx].squeeze(1), triplets['graph'][:, 1][idx].squeeze(1))), shape=(len(entity2id), len(entity2id))))
# Add transpose matrices to handle both directions of relations.
adj_list_aug = adj_list
if add_traspose_rels:
adj_list_t = [adj.T for adj in adj_list]
adj_list_aug = adj_list + adj_list_t
dgl_adj_list = ssp_multigraph_to_dgl(adj_list_aug)
return adj_list, dgl_adj_list, triplets, entity2id, relation2id, id2entity, id2relation
def intialize_worker(model, adj_list, dgl_adj_list, id2entity, params, node_features, kge_entity2id):
global model_, adj_list_, dgl_adj_list_, id2entity_, params_, node_features_, kge_entity2id_
model_, adj_list_, dgl_adj_list_, id2entity_, params_, node_features_, kge_entity2id_ = model, adj_list, dgl_adj_list, id2entity, params, node_features, kge_entity2id
def get_neg_samples_replacing_head_tail(test_links, adj_list, num_samples=50):
n, r = adj_list[0].shape[0], len(adj_list)
heads, tails, rels = test_links[:, 0], test_links[:, 1], test_links[:, 2]
neg_triplets = []
for i, (head, tail, rel) in enumerate(zip(heads, tails, rels)):
neg_triplet = {'head': [[], 0], 'tail': [[], 0]}
neg_triplet['head'][0].append([head, tail, rel])
while len(neg_triplet['head'][0]) < num_samples:
neg_head = head
neg_tail = np.random.choice(n)
if neg_head != neg_tail and adj_list[rel][neg_head, neg_tail] == 0:
neg_triplet['head'][0].append([neg_head, neg_tail, rel])
neg_triplet['tail'][0].append([head, tail, rel])
while len(neg_triplet['tail'][0]) < num_samples:
neg_head = np.random.choice(n)
neg_tail = tail
# neg_head, neg_tail, rel = np.random.choice(n), np.random.choice(n), np.random.choice(r)
if neg_head != neg_tail and adj_list[rel][neg_head, neg_tail] == 0:
neg_triplet['tail'][0].append([neg_head, neg_tail, rel])
neg_triplet['head'][0] = np.array(neg_triplet['head'][0])
neg_triplet['tail'][0] = np.array(neg_triplet['tail'][0])
neg_triplets.append(neg_triplet)
return neg_triplets
def get_neg_samples_replacing_head_tail_all(test_links, adj_list):
n, r = adj_list[0].shape[0], len(adj_list)
heads, tails, rels = test_links[:, 0], test_links[:, 1], test_links[:, 2]
neg_triplets = []
print('sampling negative triplets...')
for i, (head, tail, rel) in tqdm(enumerate(zip(heads, tails, rels)), total=len(heads)):
neg_triplet = {'head': [[], 0], 'tail': [[], 0]}
neg_triplet['head'][0].append([head, tail, rel])
for neg_tail in range(n):
neg_head = head
if neg_head != neg_tail and adj_list[rel][neg_head, neg_tail] == 0:
neg_triplet['head'][0].append([neg_head, neg_tail, rel])
neg_triplet['tail'][0].append([head, tail, rel])
for neg_head in range(n):
neg_tail = tail
if neg_head != neg_tail and adj_list[rel][neg_head, neg_tail] == 0:
neg_triplet['tail'][0].append([neg_head, neg_tail, rel])
neg_triplet['head'][0] = np.array(neg_triplet['head'][0])
neg_triplet['tail'][0] = np.array(neg_triplet['tail'][0])
neg_triplets.append(neg_triplet)
return neg_triplets
def get_neg_samples_replacing_head_tail_from_ruleN(ruleN_pred_path, entity2id, saved_relation2id):
with open(ruleN_pred_path) as f:
pred_data = [line.split() for line in f.read().split('\n')[:-1]]
neg_triplets = []
for i in range(len(pred_data) // 3):
neg_triplet = {'head': [[], 10000], 'tail': [[], 10000]}
if pred_data[3 * i][1] in saved_relation2id:
head, rel, tail = entity2id[pred_data[3 * i][0]], saved_relation2id[pred_data[3 * i][1]], entity2id[pred_data[3 * i][2]]
for j, new_head in enumerate(pred_data[3 * i + 1][1::2]):
neg_triplet['head'][0].append([entity2id[new_head], tail, rel])
if entity2id[new_head] == head:
neg_triplet['head'][1] = j
for j, new_tail in enumerate(pred_data[3 * i + 2][1::2]):
neg_triplet['tail'][0].append([head, entity2id[new_tail], rel])
if entity2id[new_tail] == tail:
neg_triplet['tail'][1] = j
neg_triplet['head'][0] = np.array(neg_triplet['head'][0])
neg_triplet['tail'][0] = np.array(neg_triplet['tail'][0])
neg_triplets.append(neg_triplet)
return neg_triplets
def incidence_matrix(adj_list):
'''
adj_list: List of sparse adjacency matrices
'''
rows, cols, dats = [], [], []
dim = adj_list[0].shape
for adj in adj_list:
adjcoo = adj.tocoo()
rows += adjcoo.row.tolist()
cols += adjcoo.col.tolist()
dats += adjcoo.data.tolist()
row = np.array(rows)
col = np.array(cols)
data = np.array(dats)
return ssp.csc_matrix((data, (row, col)), shape=dim)
def _bfs_relational(adj, roots, max_nodes_per_hop=None):
"""
BFS for graphs with multiple edge types. Returns list of level sets.
Each entry in list corresponds to relation specified by adj_list.
Modified from dgl.contrib.data.knowledge_graph to node accomodate sampling
"""
visited = set()
current_lvl = set(roots)
next_lvl = set()
while current_lvl:
for v in current_lvl:
visited.add(v)
next_lvl = _get_neighbors(adj, current_lvl)
next_lvl -= visited # set difference
if max_nodes_per_hop and max_nodes_per_hop < len(next_lvl):
next_lvl = set(random.sample(next_lvl, max_nodes_per_hop))
yield next_lvl
current_lvl = set.union(next_lvl)
def _get_neighbors(adj, nodes):
"""Takes a set of nodes and a graph adjacency matrix and returns a set of neighbors.
Directly copied from dgl.contrib.data.knowledge_graph"""
sp_nodes = _sp_row_vec_from_idx_list(list(nodes), adj.shape[1])
sp_neighbors = sp_nodes.dot(adj)
neighbors = set(ssp.find(sp_neighbors)[1]) # convert to set of indices
return neighbors
def _sp_row_vec_from_idx_list(idx_list, dim):
"""Create sparse vector of dimensionality dim from a list of indices."""
shape = (1, dim)
data = np.ones(len(idx_list))
row_ind = np.zeros(len(idx_list))
col_ind = list(idx_list)
return ssp.csr_matrix((data, (row_ind, col_ind)), shape=shape)
def get_neighbor_nodes(roots, adj, h=1, max_nodes_per_hop=None):
bfs_generator = _bfs_relational(adj, roots, max_nodes_per_hop)
lvls = list()
for _ in range(h):
try:
lvls.append(next(bfs_generator))
except StopIteration:
pass
return set().union(*lvls)
def subgraph_extraction_labeling(ind, rel, A_list, h=1, enclosing_sub_graph=False, max_nodes_per_hop=None, node_information=None, max_node_label_value=None):
# extract the h-hop enclosing subgraphs around link 'ind'
A_incidence = incidence_matrix(A_list)
A_incidence += A_incidence.T
# could pack these two into a function
root1_nei = get_neighbor_nodes(set([ind[0]]), A_incidence, h, max_nodes_per_hop)
root2_nei = get_neighbor_nodes(set([ind[1]]), A_incidence, h, max_nodes_per_hop)
subgraph_nei_nodes_int = root1_nei.intersection(root2_nei)
subgraph_nei_nodes_un = root1_nei.union(root2_nei)
# Extract subgraph | Roots being in the front is essential for labelling and the model to work properly.
if enclosing_sub_graph:
subgraph_nodes = list(ind) + list(subgraph_nei_nodes_int)
else:
subgraph_nodes = list(ind) + list(subgraph_nei_nodes_un)
subgraph = [adj[subgraph_nodes, :][:, subgraph_nodes] for adj in A_list]
labels, enclosing_subgraph_nodes = node_label_new(incidence_matrix(subgraph), max_distance=h)
pruned_subgraph_nodes = np.array(subgraph_nodes)[enclosing_subgraph_nodes].tolist()
pruned_labels = labels[enclosing_subgraph_nodes]
if max_node_label_value is not None:
pruned_labels = np.array([np.minimum(label, max_node_label_value).tolist() for label in pruned_labels])
return pruned_subgraph_nodes, pruned_labels
def remove_nodes(A_incidence, nodes):
idxs_wo_nodes = list(set(range(A_incidence.shape[1])) - set(nodes))
return A_incidence[idxs_wo_nodes, :][:, idxs_wo_nodes]
def node_label_new(subgraph, max_distance=1):
# an implementation of the proposed double-radius node labeling (DRNd L)
roots = [0, 1]
sgs_single_root = [remove_nodes(subgraph, [root]) for root in roots]
dist_to_roots = [np.clip(ssp.csgraph.dijkstra(sg, indices=[0], directed=False, unweighted=True, limit=1e6)[:, 1:], 0, 1e7) for r, sg in enumerate(sgs_single_root)]
dist_to_roots = np.array(list(zip(dist_to_roots[0][0], dist_to_roots[1][0])), dtype=int)
# dist_to_roots[np.abs(dist_to_roots) > 1e6] = 0
# dist_to_roots = dist_to_roots + 1
target_node_labels = np.array([[0, 1], [1, 0]])
labels = np.concatenate((target_node_labels, dist_to_roots)) if dist_to_roots.size else target_node_labels
enclosing_subgraph_nodes = np.where(np.max(labels, axis=1) <= max_distance)[0]
# print(len(enclosing_subgraph_nodes))
return labels, enclosing_subgraph_nodes
def ssp_multigraph_to_dgl(graph, n_feats=None):
"""
Converting ssp multigraph (i.e. list of adjs) to dgl multigraph.
"""
g_nx = nx.MultiDiGraph()
g_nx.add_nodes_from(list(range(graph[0].shape[0])))
# Add edges
for rel, adj in enumerate(graph):
# Convert adjacency matrix to tuples for nx0
nx_triplets = []
for src, dst in list(zip(adj.tocoo().row, adj.tocoo().col)):
nx_triplets.append((src, dst, {'type': rel}))
g_nx.add_edges_from(nx_triplets)
# make dgl graph
g_dgl = dgl.DGLGraph(multigraph=True)
g_dgl.from_networkx(g_nx, edge_attrs=['type'])
# add node features
if n_feats is not None:
g_dgl.ndata['feat'] = torch.tensor(n_feats)
return g_dgl
def prepare_features(subgraph, n_labels, max_n_label, n_feats=None):
# One hot encode the node label feature and concat to n_featsure
n_nodes = subgraph.number_of_nodes()
label_feats = np.zeros((n_nodes, max_n_label[0] + 1 + max_n_label[1] + 1))
label_feats[np.arange(n_nodes), n_labels[:, 0]] = 1
label_feats[np.arange(n_nodes), max_n_label[0] + 1 + n_labels[:, 1]] = 1
n_feats = np.concatenate((label_feats, n_feats), axis=1) if n_feats is not None else label_feats
subgraph.ndata['feat'] = torch.FloatTensor(n_feats)
head_id = np.argwhere([label[0] == 0 and label[1] == 1 for label in n_labels])
tail_id = np.argwhere([label[0] == 1 and label[1] == 0 for label in n_labels])
n_ids = np.zeros(n_nodes)
n_ids[head_id] = 1 # head
n_ids[tail_id] = 2 # tail
subgraph.ndata['id'] = torch.FloatTensor(n_ids)
return subgraph
def get_subgraphs(all_links, adj_list, dgl_adj_list, max_node_label_value, id2entity, node_features=None, kge_entity2id=None):
# dgl_adj_list = ssp_multigraph_to_dgl(adj_list)
subgraphs = []
r_labels = []
for link in all_links:
head, tail, rel = link[0], link[1], link[2]
nodes, node_labels = subgraph_extraction_labeling((head, tail), rel, adj_list, h=params_.hop, enclosing_sub_graph=params.enclosing_sub_graph, max_node_label_value=max_node_label_value)
subgraph = dgl.DGLGraph(dgl_adj_list.subgraph(nodes))
subgraph.edata['type'] = dgl_adj_list.edata['type'][dgl_adj_list.subgraph(nodes).parent_eid]
subgraph.edata['label'] = torch.tensor(rel * np.ones(subgraph.edata['type'].shape), dtype=torch.long)
edges_btw_roots = subgraph.edge_id(0, 1)
rel_link = np.nonzero(subgraph.edata['type'][edges_btw_roots] == rel)
if rel_link.squeeze().nelement() == 0:
# subgraph.add_edge(0, 1, {'type': torch.tensor([rel]), 'label': torch.tensor([rel])})
subgraph.add_edge(0, 1)
subgraph.edata['type'][-1] = torch.tensor(rel).type(torch.LongTensor)
subgraph.edata['label'][-1] = torch.tensor(rel).type(torch.LongTensor)
kge_nodes = [kge_entity2id[id2entity[n]] for n in nodes] if kge_entity2id else None
n_feats = node_features[kge_nodes] if node_features is not None else None
subgraph = prepare_features(subgraph, node_labels, max_node_label_value, n_feats)
subgraphs.append(subgraph)
r_labels.append(rel)
batched_graph = dgl.batch(subgraphs)
r_labels = torch.LongTensor(r_labels)
return (batched_graph, r_labels)
def get_rank(neg_links):
head_neg_links = neg_links['head'][0]
head_target_id = neg_links['head'][1]
if head_target_id != 10000:
data = get_subgraphs(head_neg_links, adj_list_, dgl_adj_list_, model_.gnn.max_label_value, id2entity_, node_features_, kge_entity2id_)
head_scores = model_(data).squeeze(1).detach().numpy()
head_rank = np.argwhere(np.argsort(head_scores)[::-1] == head_target_id) + 1
else:
head_scores = np.array([])
head_rank = 10000
tail_neg_links = neg_links['tail'][0]
tail_target_id = neg_links['tail'][1]
if tail_target_id != 10000:
data = get_subgraphs(tail_neg_links, adj_list_, dgl_adj_list_, model_.gnn.max_label_value, id2entity_, node_features_, kge_entity2id_)
tail_scores = model_(data).squeeze(1).detach().numpy()
tail_rank = np.argwhere(np.argsort(tail_scores)[::-1] == tail_target_id) + 1
else:
tail_scores = np.array([])
tail_rank = 10000
return head_scores, head_rank, tail_scores, tail_rank
def save_to_file(neg_triplets, id2entity, id2relation):
with open(os.path.join('./data', params.dataset, 'ranking_head.txt'), "w") as f:
for neg_triplet in neg_triplets:
for s, o, r in neg_triplet['head'][0]:
f.write('\t'.join([id2entity[s], id2relation[r], id2entity[o]]) + '\n')
with open(os.path.join('./data', params.dataset, 'ranking_tail.txt'), "w") as f:
for neg_triplet in neg_triplets:
for s, o, r in neg_triplet['tail'][0]:
f.write('\t'.join([id2entity[s], id2relation[r], id2entity[o]]) + '\n')
def save_score_to_file(neg_triplets, all_head_scores, all_tail_scores, id2entity, id2relation):
with open(os.path.join('./data', params.dataset, 'grail_ranking_head_predictions.txt'), "w") as f:
for i, neg_triplet in enumerate(neg_triplets):
for [s, o, r], head_score in zip(neg_triplet['head'][0], all_head_scores[50 * i:50 * (i + 1)]):
f.write('\t'.join([id2entity[s], id2relation[r], id2entity[o], str(head_score)]) + '\n')
with open(os.path.join('./data', params.dataset, 'grail_ranking_tail_predictions.txt'), "w") as f:
for i, neg_triplet in enumerate(neg_triplets):
for [s, o, r], tail_score in zip(neg_triplet['tail'][0], all_tail_scores[50 * i:50 * (i + 1)]):
f.write('\t'.join([id2entity[s], id2relation[r], id2entity[o], str(tail_score)]) + '\n')
def save_score_to_file_from_ruleN(neg_triplets, all_head_scores, all_tail_scores, id2entity, id2relation):
with open(os.path.join('./data', params.dataset, 'grail_ruleN_ranking_head_predictions.txt'), "w") as f:
for i, neg_triplet in enumerate(neg_triplets):
for [s, o, r], head_score in zip(neg_triplet['head'][0], all_head_scores[50 * i:50 * (i + 1)]):
f.write('\t'.join([id2entity[s], id2relation[r], id2entity[o], str(head_score)]) + '\n')
with open(os.path.join('./data', params.dataset, 'grail_ruleN_ranking_tail_predictions.txt'), "w") as f:
for i, neg_triplet in enumerate(neg_triplets):
for [s, o, r], tail_score in zip(neg_triplet['tail'][0], all_tail_scores[50 * i:50 * (i + 1)]):
f.write('\t'.join([id2entity[s], id2relation[r], id2entity[o], str(tail_score)]) + '\n')
def get_kge_embeddings(dataset, kge_model):
path = './experiments/kge_baselines/{}_{}'.format(kge_model, dataset)
node_features = np.load(os.path.join(path, 'entity_embedding.npy'))
with open(os.path.join(path, 'id2entity.json')) as json_file:
kge_id2entity = json.load(json_file)
kge_entity2id = {v: int(k) for k, v in kge_id2entity.items()}
return node_features, kge_entity2id
def main(params):
model = torch.load(params.model_path, map_location='cpu')
adj_list, dgl_adj_list, triplets, entity2id, relation2id, id2entity, id2relation = process_files(params.file_paths, model.relation2id, params.add_traspose_rels)
node_features, kge_entity2id = get_kge_embeddings(params.dataset, params.kge_model) if params.use_kge_embeddings else (None, None)
if params.mode == 'sample':
neg_triplets = get_neg_samples_replacing_head_tail(triplets['links'], adj_list)
save_to_file(neg_triplets, id2entity, id2relation)
elif params.mode == 'all':
neg_triplets = get_neg_samples_replacing_head_tail_all(triplets['links'], adj_list)
elif params.mode == 'ruleN':
neg_triplets = get_neg_samples_replacing_head_tail_from_ruleN(params.ruleN_pred_path, entity2id, relation2id)
ranks = []
all_head_scores = []
all_tail_scores = []
with mp.Pool(processes=None, initializer=intialize_worker, initargs=(model, adj_list, dgl_adj_list, id2entity, params, node_features, kge_entity2id)) as p:
for head_scores, head_rank, tail_scores, tail_rank in tqdm(p.imap(get_rank, neg_triplets), total=len(neg_triplets)):
ranks.append(head_rank)
ranks.append(tail_rank)
all_head_scores += head_scores.tolist()
all_tail_scores += tail_scores.tolist()
# intialize_worker(model, adj_list, dgl_adj_list, id2entity, params, node_features, kge_entity2id)
# for link in tqdm(neg_triplets, total=len(neg_triplets)):
# head_scores, head_rank, tail_scores, tail_rank = get_rank(link)
# ranks.append(head_rank)
# ranks.append(tail_rank)
# all_head_scores += head_scores.tolist()
# all_tail_scores += tail_scores.tolist()
if params.mode == 'ruleN':
save_score_to_file_from_ruleN(neg_triplets, all_head_scores, all_tail_scores, id2entity, id2relation)
else:
save_score_to_file(neg_triplets, all_head_scores, all_tail_scores, id2entity, id2relation)
isHit1List = [x for x in ranks if x <= 1]
isHit5List = [x for x in ranks if x <= 5]
isHit10List = [x for x in ranks if x <= 10]
hits_1 = len(isHit1List) / len(ranks)
hits_5 = len(isHit5List) / len(ranks)
hits_10 = len(isHit10List) / len(ranks)
mrr = np.mean(1 / np.array(ranks))
logger.info(f'MRR | Hits@1 | Hits@5 | Hits@10 : {mrr} | {hits_1} | {hits_5} | {hits_10}')
if __name__ == '__main__':
logging.basicConfig(level=logging.INFO)
parser = argparse.ArgumentParser(description='Testing script for hits@10')
# Experiment setup params
parser.add_argument("--experiment_name", "-e", type=str, default="fb_v2_margin_loss",
help="Experiment name. Log file with this name will be created")
parser.add_argument("--dataset", "-d", type=str, default="FB237_v2",
help="Path to dataset")
parser.add_argument("--mode", "-m", type=str, default="sample", choices=["sample", "all", "ruleN"],
help="Negative sampling mode")
parser.add_argument("--use_kge_embeddings", "-kge", type=bool, default=False,
help='whether to use pretrained KGE embeddings')
parser.add_argument("--kge_model", type=str, default="TransE",
help="Which KGE model to load entity embeddings from")
parser.add_argument('--enclosing_sub_graph', '-en', type=bool, default=True,
help='whether to only consider enclosing subgraph')
parser.add_argument("--hop", type=int, default=3,
help="How many hops to go while eextracting subgraphs?")
parser.add_argument('--add_traspose_rels', '-tr', type=bool, default=False,
help='Whether to append adj matrix list with symmetric relations?')
params = parser.parse_args()
params.file_paths = {
'graph': os.path.join('./data', params.dataset, 'train.txt'),
'links': os.path.join('./data', params.dataset, 'test.txt')
}
params.ruleN_pred_path = os.path.join('./data', params.dataset, 'pos_predictions.txt')
params.model_path = os.path.join('experiments', params.experiment_name, 'best_graph_classifier.pth')
file_handler = logging.FileHandler(os.path.join('experiments', params.experiment_name, f'log_rank_test_{time.time()}.txt'))
logger = logging.getLogger()
logger.addHandler(file_handler)
logger.info('============ Initialized logger ============')
logger.info('\n'.join('%s: %s' % (k, str(v)) for k, v
in sorted(dict(vars(params)).items())))
logger.info('============================================')
main(params)