-
Notifications
You must be signed in to change notification settings - Fork 0
/
utils.py
89 lines (73 loc) · 2.15 KB
/
utils.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
import numpy as np
from scipy import spatial
import networkx as nx
import os
import pickle
def process_title(title):
return [w.replace('(', '').replace(')','').lower() for w in title.split('_')]
'''
Data
'''
data_path = 'data'
def load_graph():
'''
Loads the Wikipedia subgraph and the translation dictionaries.
@Returns:
graph... networkx graph
idx2word... dictionary (key=index, value=word)
word2idx... dictionary (key=word, value=idx)
'''
graph = nx.read_gpickle(os.path.join(data_path, 'wales', 'wiki_graph.p'))
with open(os.path.join(data_path, 'wales', 'idx2name.p'), 'rb') as handle:
idx2word = pickle.load(handle)
with open(os.path.join(data_path, 'wales', 'name2idx.p'), 'rb') as handle:
word2idx = pickle.load(handle)
return graph, idx2word, word2idx
def load_challenges():
'''
Loads scraped Wikirace challenges.
@Returns:
challenges... list of word tuples
'''
challenges = []
with open(os.path.join(data_path, 'wales', 'challenges.txt'), 'r') as inf:
lines = inf.readlines()
for line in lines:
s, t = line.replace('\n', '').split('\t')
challenges.append((s,t))
return challenges
def load_embedding(emb_file):
with open(os.path.join('data', 'embeddings', emb_file), 'rb') as handle:
emb_dict = pickle.load(handle)
return emb_dict
'''
Distance Functions
'''
def cosine_dist(v1, v2):
'''
Cosine distance between two vectors.
@Params:
v1... vector 1
v2... vector 2
@Returns:
1-cossim
'''
d = spatial.distance.cosine(v1, v2)
if np.isnan(d) or np.all(v1==0) or np.all(v2==0):
return np.inf
else:
return d
def cosine_dist_parallel(V, v):
'''
Parallel cosine distance between batch of vectors and a vector.
@Params:
V... batch of vectors (n x d)
v... vector (d)
@Returns:
1-cossim between each vector in batch and v
'''
mask = np.sum((V!=0), axis=1)==0
tmp = V@v
tmp[mask] = np.inf
tmp[~mask] = 1 - tmp[~mask]/(np.linalg.norm(V[~mask,:], axis=1) * np.linalg.norm(v))
return tmp