aboutsummaryrefslogtreecommitdiffstats
path: root/WMD_retrieval.py
diff options
context:
space:
mode:
Diffstat (limited to 'WMD_retrieval.py')
-rw-r--r--WMD_retrieval.py334
1 files changed, 100 insertions, 234 deletions
diff --git a/WMD_retrieval.py b/WMD_retrieval.py
index f99eaa1..b49ba7d 100644
--- a/WMD_retrieval.py
+++ b/WMD_retrieval.py
@@ -1,259 +1,125 @@
1###########################
2# Wasserstein Retrieval #
3###########################
4import argparse 1import argparse
5
6parser = argparse.ArgumentParser(description='run retrieval using wmd and wasserstein distances')
7parser.add_argument('source_lang', help='source language short name')
8parser.add_argument('target_lang', help='target language short name')
9parser.add_argument('source_vector', help='path of the source vector')
10parser.add_argument('target_vector', help='path of the target vector')
11parser.add_argument('source_defs', help='path of the source definitions')
12parser.add_argument('target_defs', help='path of the target definitions')
13parser.add_argument('-n', '--instances', help='number of instances in each language to retrieve', default=2000, type=int)
14
15args = parser.parse_args()
16
17source_lang = args.source_lang
18target_lang = args.target_lang
19
20def load_embeddings(path, dimension=300):
21 """
22 Loads the embeddings from a word2vec formatted file.
23 word2vec format is one line per word and it's associated embedding
24 (dimension x floating numbers) separated by spaces
25 The first line may or may not include the word count and dimension
26 """
27 vectors = {}
28 with open(path, mode='r', encoding='utf8') as fp:
29 first_line = fp.readline().rstrip('\n')
30 if first_line.count(' ') == 1:
31 # includes the "word_count dimension" information
32 (word_count, dimension) = map(int, first_line.split())
33 else: # assume the file only contains vectors
34 fp.seek(0)
35 for line in fp:
36 elems = line.split()
37 vectors[" ".join(elems[:-dimension])] = " ".join(elems[-dimension:])
38 return vectors
39
40#######################################################################
41# Vectors Load Here #
42#######################################################################
43
44source_vectors_filename = args.source_vector
45target_vectors_filename = args.target_vector
46vectors_source = load_embeddings(source_vectors_filename)
47vectors_target = load_embeddings(target_vectors_filename)
48
49#######################################################################
50# Corpora Load Here #
51#######################################################################
52source_defs_filename = args.source_defs
53target_defs_filename = args.target_defs
54defs_source = [line.rstrip('\n') for line in open(source_defs_filename, encoding='utf8')]
55defs_target = [line.rstrip('\n') for line in open(target_defs_filename, encoding='utf8')]
56
57import numpy as np 2import numpy as np
58from mosestokenizer import *
59
60def clean_corpus_using_embeddings_vocabulary(
61 embeddings_dictionary,
62 corpus,
63 vectors,
64 language,
65 ):
66 '''
67 Cleans corpus using the dictionary of embeddings.
68 Any word without an associated embedding in the dictionary is ignored.
69 Adds '__target-language' and '__source-language' at the end of the words according to their language.
70 '''
71 clean_corpus, clean_vectors, keys = [], {}, []
72 words_we_want = set(embeddings_dictionary)
73 tokenize = MosesTokenizer(language)
74 for key, doc in enumerate(corpus):
75 clean_doc = []
76 words = tokenize(doc)
77 for word in words:
78 if word in words_we_want:
79 clean_doc.append(word + '__%s' % language)
80 clean_vectors[word + '__%s' % language] = np.array(vectors[word].split()).astype(np.float)
81 if len(clean_doc) > 3 and len(clean_doc) < 25:
82 keys.append(key)
83 clean_corpus.append(' '.join(clean_doc))
84 tokenize.close()
85 return np.array(clean_corpus), clean_vectors, keys
86
87clean_src_corpus, clean_src_vectors, src_keys = clean_corpus_using_embeddings_vocabulary(
88 set(vectors_source.keys()),
89 defs_source,
90 vectors_source,
91 source_lang,
92 )
93
94clean_target_corpus, clean_target_vectors, target_keys = clean_corpus_using_embeddings_vocabulary(
95 set(vectors_target.keys()),
96 defs_target,
97 vectors_target,
98 target_lang,
99 )
100
101# Here is the part Wasserstein prunes two corporas to 500 articles each
102# Our dataset does not have that luxury (turns out it's not a luxury but a necessity)
103
104import random 3import random
105take = args.instances 4from sklearn.feature_extraction.text import CountVectorizer, TfidfVectorizer
5from sklearn.preprocessing import normalize
6from Wasserstein_Distance import Wasserstein_Retriever
7from Wasserstein_Distance import load_embeddings, clean_corpus_using_embeddings_vocabulary, mrr_precision_at_k
8import csv
9import sys
106 10
107common_keys = set(src_keys).intersection(set(target_keys)) 11def main(args):
108take = min(len(common_keys), take) # you can't sample more than length
109experiment_keys = random.sample(common_keys, take)
110 12
111instances = len(experiment_keys) 13 np.seterr(divide='ignore') # POT has issues with divide by zero errors
14 source_lang = args.source_lang
15 target_lang = args.target_lang
112 16
113clean_src_corpus = list(clean_src_corpus[experiment_keys]) 17 source_vectors_filename = args.source_vector
114clean_target_corpus = list(clean_target_corpus[experiment_keys]) 18 target_vectors_filename = args.target_vector
19 vectors_source = load_embeddings(source_vectors_filename)
20 vectors_target = load_embeddings(target_vectors_filename)
115 21
116print(f'{source_lang} - {target_lang} : document sizes: {len(clean_src_corpus)}, {len(clean_target_corpus)}') 22 source_defs_filename = args.source_defs
23 target_defs_filename = args.target_defs
117 24
118del vectors_source, vectors_target, defs_source, defs_target 25 batch = args.batch
26 mode = args.mode
27 runfor = list()
119 28
120from sklearn.feature_extraction.text import CountVectorizer, TfidfVectorizer 29 if (mode == 'all'):
121 30 runfor.extend(['wmd','snk'])
122vec = CountVectorizer().fit(clean_src_corpus + clean_target_corpus)
123common = [word for word in vec.get_feature_names() if word in clean_src_vectors or word in clean_target_vectors]
124W_common = []
125for w in common:
126 if w in clean_src_vectors:
127 W_common.append(np.array(clean_src_vectors[w]))
128 else: 31 else:
129 W_common.append(np.array(clean_target_vectors[w])) 32 runfor.append(mode)
130 33
131print(f'{source_lang} - {target_lang}: the vocabulary size is {len(W_common)}') 34 defs_source = [line.rstrip('\n') for line in open(source_defs_filename, encoding='utf8')]
35 defs_target = [line.rstrip('\n') for line in open(target_defs_filename, encoding='utf8')]
132 36
133from sklearn.preprocessing import normalize 37 clean_src_corpus, clean_src_vectors, src_keys = clean_corpus_using_embeddings_vocabulary(
134W_common = np.array(W_common) 38 set(vectors_source.keys()),
135W_common = normalize(W_common) 39 defs_source,
136vect = TfidfVectorizer(vocabulary=common, dtype=np.double, norm=None) 40 vectors_source,
137vect.fit(clean_src_corpus + clean_target_corpus) 41 source_lang,
138X_train_idf = vect.transform(clean_src_corpus) 42 )
139X_test_idf = vect.transform(clean_target_corpus)
140 43
141vect_tf = CountVectorizer(vocabulary=common, dtype=np.double) 44 clean_target_corpus, clean_target_vectors, target_keys = clean_corpus_using_embeddings_vocabulary(
142vect_tf.fit(clean_src_corpus + clean_target_corpus) 45 set(vectors_target.keys()),
143X_train_tf = vect_tf.transform(clean_src_corpus) 46 defs_target,
144X_test_tf = vect_tf.transform(clean_target_corpus) 47 vectors_target,
48 target_lang,
49 )
145 50
146import ot 51 take = args.instances
147from sklearn.neighbors import KNeighborsClassifier
148from sklearn.metrics import euclidean_distances
149from sklearn.externals.joblib import Parallel, delayed
150from sklearn.utils import check_array
151from sklearn.metrics.scorer import check_scoring
152from pathos.multiprocessing import ProcessingPool as Pool
153from sklearn.metrics import euclidean_distances
154 52
155class WassersteinDistances(KNeighborsClassifier): 53 common_keys = set(src_keys).intersection(set(target_keys))
156 """ 54 take = min(len(common_keys), take) # you can't sample more than length
157 Implements a nearest neighbors classifier for input distributions using the Wasserstein distance as metric. 55 experiment_keys = random.sample(common_keys, take)
158 Source and target distributions are l_1 normalized before computing the Wasserstein distance.
159 Wasserstein is parametrized by the distances between the individual points of the distributions.
160 In this work, we propose to use cross-lingual embeddings for calculating these distances.
161 56
162 """ 57 instances = len(experiment_keys)
163 def __init__(self, W_embed, n_neighbors=1, n_jobs=1, verbose=False, sinkhorn= False, sinkhorn_reg=0.1):
164 """
165 Initialization of the class.
166 Arguments
167 ---------
168 W_embed: embeddings of the words, np.array
169 verbose: True/False
170 """
171 self.sinkhorn = sinkhorn
172 self.sinkhorn_reg = sinkhorn_reg
173 self.W_embed = W_embed
174 self.verbose = verbose
175 super(WassersteinDistances, self).__init__(n_neighbors=n_neighbors, n_jobs=n_jobs, metric='precomputed', algorithm='brute')
176 58
177 def _wmd(self, i, row, X_train): 59 clean_src_corpus = list(clean_src_corpus[experiment_keys])
178 union_idx = np.union1d(X_train[i].indices, row.indices) 60 clean_target_corpus = list(clean_target_corpus[experiment_keys])
179 W_minimal = self.W_embed[union_idx]
180 W_dist = euclidean_distances(W_minimal)
181 bow_i = X_train[i, union_idx].A.ravel()
182 bow_j = row[:, union_idx].A.ravel()
183 if self.sinkhorn:
184 return ot.sinkhorn2(bow_i, bow_j, W_dist, self.sinkhorn_reg, numItermax=50, method='sinkhorn_stabilized',)[0]
185 else:
186 return ot.emd2(bow_i, bow_j, W_dist)
187 61
188 def _wmd_row(self, row): 62 if (not batch):
189 X_train = self._fit_X 63 print(f'{source_lang} - {target_lang} : document sizes: {len(clean_src_corpus)}, {len(clean_target_corpus)}')
190 n_samples_train = X_train.shape[0]
191 return [self._wmd(i, row, X_train) for i in range(n_samples_train)]
192 64
193 def _pairwise_wmd(self, X_test, X_train=None): 65 del vectors_source, vectors_target, defs_source, defs_target
194 n_samples_test = X_test.shape[0]
195 66
196 if X_train is None: 67 vec = CountVectorizer().fit(clean_src_corpus + clean_target_corpus)
197 X_train = self._fit_X 68 common = [word for word in vec.get_feature_names() if word in clean_src_vectors or word in clean_target_vectors]
198 pool = Pool(nodes=self.n_jobs) # Parallelization of the calculation of the distances 69 W_common = []
199 dist = pool.map(self._wmd_row, X_test) 70 for w in common:
200 return np.array(dist) 71 if w in clean_src_vectors:
201 72 W_common.append(np.array(clean_src_vectors[w]))
202 def fit(self, X, y): 73 else:
203 X = check_array(X, accept_sparse='csr', copy=True) 74 W_common.append(np.array(clean_target_vectors[w]))
204 X = normalize(X, norm='l1', copy=False) 75
205 return super(WassersteinDistances, self).fit(X, y) 76 if (not batch):
206 77 print(f'{source_lang} - {target_lang}: the vocabulary size is {len(W_common)}')
207 def predict(self, X): 78
208 X = check_array(X, accept_sparse='csr', copy=True) 79 W_common = np.array(W_common)
209 X = normalize(X, norm='l1', copy=False) 80 W_common = normalize(W_common)
210 dist = self._pairwise_wmd(X) 81 vect = TfidfVectorizer(vocabulary=common, dtype=np.double, norm=None)
211 return super(WassersteinDistances, self).predict(dist) 82 vect.fit(clean_src_corpus + clean_target_corpus)
212 83 X_train_idf = vect.transform(clean_src_corpus)
213 def kneighbors(self, X, n_neighbors=1): 84 X_test_idf = vect.transform(clean_target_corpus)
214 X = check_array(X, accept_sparse='csr', copy=True) 85
215 X = normalize(X, norm='l1', copy=False) 86 vect_tf = CountVectorizer(vocabulary=common, dtype=np.double)
216 dist = self._pairwise_wmd(X) 87 vect_tf.fit(clean_src_corpus + clean_target_corpus)
217 return super(WassersteinDistances, self).kneighbors(dist, n_neighbors) 88 X_train_tf = vect_tf.transform(clean_src_corpus)
89 X_test_tf = vect_tf.transform(clean_target_corpus)
90
91 for metric in runfor:
92 if (not batch):
93 print(f'{metric} - tfidf: {source_lang} - {target_lang}')
94
95 clf = WassersteinDistances(W_embed=W_common, n_neighbors=5, n_jobs=14, sinkhorn=(metric == 'snk'))
96 clf.fit(X_train_idf[:instances], np.ones(instances))
97 dist, preds = clf.kneighbors(X_test_idf[:instances], n_neighbors=instances)
98 mrr, p_at_1 = mrr_precision_at_k(list(range(len(preds))), preds)
99 percentage = p_at_1 * 100
100
101 if (not batch):
102 print(f'MRR: {mrr} | Precision @ 1: {p_at_1}')
103 else:
104 fields = [f'{source_lang}', f'{target_lang}', f'{instances}', f'{mrr}', f'{p_at_1}', f'{percentage}']
105 with open(f'{metric}_retrieval_result.csv', 'a') as f:
106 writer = csv.writer(f)
107 writer.writerow(fields)
218 108
219def mrr_precision_at_k(golden, preds, k_list=[1,]):
220 """
221 Calculates Mean Reciprocal Error and Hits@1 == Precision@1
222 """
223 my_score = 0
224 precision_at = np.zeros(len(k_list))
225 for key, elem in enumerate(golden):
226 if elem in preds[key]:
227 location = np.where(preds[key]==elem)[0][0]
228 my_score += 1/(1+ location)
229 for k_index, k_value in enumerate(k_list):
230 if location < k_value:
231 precision_at[k_index] += 1
232 return my_score/len(golden), (precision_at/len(golden))[0]
233 109
234print(f'WMD - tfidf: {source_lang} - {target_lang}') 110if __name__ == "__main__":
235clf = WassersteinDistances(W_embed=W_common, n_neighbors=5, n_jobs=14)
236clf.fit(X_train_idf[:instances], np.ones(instances))
237dist, preds = clf.kneighbors(X_test_idf[:instances], n_neighbors=instances)
238mrr, p_at_1 = mrr_precision_at_k(list(range(len(preds))), preds)
239print(f'MRR: {mrr} | Precision @ 1: {p_at_1}')
240 111
241import csv 112 parser = argparse.ArgumentParser(description='run retrieval using wmd or snk')
242percentage = p_at_1 * 100 113 parser.add_argument('source_lang', help='source language short name')
243fields = [f'{source_lang}', f'{target_lang}', f'{instances}', f'{mrr}', f'{p_at_1}', f'{percentage}'] 114 parser.add_argument('target_lang', help='target language short name')
244with open('/home/syigit/multilang_results/wmd_retrieval_result.csv', 'a') as f: 115 parser.add_argument('source_vector', help='path of the source vector')
245 writer = csv.writer(f) 116 parser.add_argument('target_vector', help='path of the target vector')
246 writer.writerow(fields) 117 parser.add_argument('source_defs', help='path of the source definitions')
118 parser.add_argument('target_defs', help='path of the target definitions')
119 parser.add_argument('-b', '--batch', action='store_true', help='running in batch (store results in csv) or running a single instance (output the results)')
120 parser.add_argument('mode', choices=['all', 'wmd', 'snk'], default='all', help='which methods to run')
121 parser.add_argument('-n', '--instances', help='number of instances in each language to retrieve', default=2000, type=int)
247 122
248print(f'Sinkhorn - tfidf: {source_lang} - {target_lang}') 123 args = parser.parse_args()
249clf = WassersteinDistances(W_embed=W_common, n_neighbors=5, n_jobs=14, sinkhorn=True)
250clf.fit(X_train_idf[:instances], np.ones(instances))
251dist, preds = clf.kneighbors(X_test_idf[:instances], n_neighbors=instances)
252mrr, p_at_1 = mrr_precision_at_k(list(range(len(preds))), preds)
253print(f'MRR: {mrr} | Precision @ 1: {p_at_1}')
254 124
255percentage = p_at_1 * 100 125 main(args)
256fields = [f'{source_lang}', f'{target_lang}', f'{instances}', f'{mrr}', f'{p_at_1}', f'{percentage}']
257with open('/home/syigit/multilang_results/sinkhorn_retrieval_result.csv', 'a') as f:
258 writer = csv.writer(f)
259 writer.writerow(fields)