diff options
Diffstat (limited to 'WMD_retrieval.py')
-rw-r--r-- | WMD_retrieval.py | 334 |
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 | ########################### | ||
4 | import argparse | 1 | import argparse |
5 | |||
6 | parser = argparse.ArgumentParser(description='run retrieval using wmd and wasserstein distances') | ||
7 | parser.add_argument('source_lang', help='source language short name') | ||
8 | parser.add_argument('target_lang', help='target language short name') | ||
9 | parser.add_argument('source_vector', help='path of the source vector') | ||
10 | parser.add_argument('target_vector', help='path of the target vector') | ||
11 | parser.add_argument('source_defs', help='path of the source definitions') | ||
12 | parser.add_argument('target_defs', help='path of the target definitions') | ||
13 | parser.add_argument('-n', '--instances', help='number of instances in each language to retrieve', default=2000, type=int) | ||
14 | |||
15 | args = parser.parse_args() | ||
16 | |||
17 | source_lang = args.source_lang | ||
18 | target_lang = args.target_lang | ||
19 | |||
20 | def 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 | |||
44 | source_vectors_filename = args.source_vector | ||
45 | target_vectors_filename = args.target_vector | ||
46 | vectors_source = load_embeddings(source_vectors_filename) | ||
47 | vectors_target = load_embeddings(target_vectors_filename) | ||
48 | |||
49 | ####################################################################### | ||
50 | # Corpora Load Here # | ||
51 | ####################################################################### | ||
52 | source_defs_filename = args.source_defs | ||
53 | target_defs_filename = args.target_defs | ||
54 | defs_source = [line.rstrip('\n') for line in open(source_defs_filename, encoding='utf8')] | ||
55 | defs_target = [line.rstrip('\n') for line in open(target_defs_filename, encoding='utf8')] | ||
56 | |||
57 | import numpy as np | 2 | import numpy as np |
58 | from mosestokenizer import * | ||
59 | |||
60 | def 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 | |||
87 | clean_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 | |||
94 | clean_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 | |||
104 | import random | 3 | import random |
105 | take = args.instances | 4 | from sklearn.feature_extraction.text import CountVectorizer, TfidfVectorizer |
5 | from sklearn.preprocessing import normalize | ||
6 | from Wasserstein_Distance import Wasserstein_Retriever | ||
7 | from Wasserstein_Distance import load_embeddings, clean_corpus_using_embeddings_vocabulary, mrr_precision_at_k | ||
8 | import csv | ||
9 | import sys | ||
106 | 10 | ||
107 | common_keys = set(src_keys).intersection(set(target_keys)) | 11 | def main(args): |
108 | take = min(len(common_keys), take) # you can't sample more than length | ||
109 | experiment_keys = random.sample(common_keys, take) | ||
110 | 12 | ||
111 | instances = 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 | ||
113 | clean_src_corpus = list(clean_src_corpus[experiment_keys]) | 17 | source_vectors_filename = args.source_vector |
114 | clean_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 | ||
116 | print(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 | ||
118 | del vectors_source, vectors_target, defs_source, defs_target | 25 | batch = args.batch |
26 | mode = args.mode | ||
27 | runfor = list() | ||
119 | 28 | ||
120 | from sklearn.feature_extraction.text import CountVectorizer, TfidfVectorizer | 29 | if (mode == 'all'): |
121 | 30 | runfor.extend(['wmd','snk']) | |
122 | vec = CountVectorizer().fit(clean_src_corpus + clean_target_corpus) | ||
123 | common = [word for word in vec.get_feature_names() if word in clean_src_vectors or word in clean_target_vectors] | ||
124 | W_common = [] | ||
125 | for 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 | ||
131 | print(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 | ||
133 | from sklearn.preprocessing import normalize | 37 | clean_src_corpus, clean_src_vectors, src_keys = clean_corpus_using_embeddings_vocabulary( |
134 | W_common = np.array(W_common) | 38 | set(vectors_source.keys()), |
135 | W_common = normalize(W_common) | 39 | defs_source, |
136 | vect = TfidfVectorizer(vocabulary=common, dtype=np.double, norm=None) | 40 | vectors_source, |
137 | vect.fit(clean_src_corpus + clean_target_corpus) | 41 | source_lang, |
138 | X_train_idf = vect.transform(clean_src_corpus) | 42 | ) |
139 | X_test_idf = vect.transform(clean_target_corpus) | ||
140 | 43 | ||
141 | vect_tf = CountVectorizer(vocabulary=common, dtype=np.double) | 44 | clean_target_corpus, clean_target_vectors, target_keys = clean_corpus_using_embeddings_vocabulary( |
142 | vect_tf.fit(clean_src_corpus + clean_target_corpus) | 45 | set(vectors_target.keys()), |
143 | X_train_tf = vect_tf.transform(clean_src_corpus) | 46 | defs_target, |
144 | X_test_tf = vect_tf.transform(clean_target_corpus) | 47 | vectors_target, |
48 | target_lang, | ||
49 | ) | ||
145 | 50 | ||
146 | import ot | 51 | take = args.instances |
147 | from sklearn.neighbors import KNeighborsClassifier | ||
148 | from sklearn.metrics import euclidean_distances | ||
149 | from sklearn.externals.joblib import Parallel, delayed | ||
150 | from sklearn.utils import check_array | ||
151 | from sklearn.metrics.scorer import check_scoring | ||
152 | from pathos.multiprocessing import ProcessingPool as Pool | ||
153 | from sklearn.metrics import euclidean_distances | ||
154 | 52 | ||
155 | class 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 | ||
219 | def 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 | ||
234 | print(f'WMD - tfidf: {source_lang} - {target_lang}') | 110 | if __name__ == "__main__": |
235 | clf = WassersteinDistances(W_embed=W_common, n_neighbors=5, n_jobs=14) | ||
236 | clf.fit(X_train_idf[:instances], np.ones(instances)) | ||
237 | dist, preds = clf.kneighbors(X_test_idf[:instances], n_neighbors=instances) | ||
238 | mrr, p_at_1 = mrr_precision_at_k(list(range(len(preds))), preds) | ||
239 | print(f'MRR: {mrr} | Precision @ 1: {p_at_1}') | ||
240 | 111 | ||
241 | import csv | 112 | parser = argparse.ArgumentParser(description='run retrieval using wmd or snk') |
242 | percentage = p_at_1 * 100 | 113 | parser.add_argument('source_lang', help='source language short name') |
243 | fields = [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') |
244 | with 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 | ||
248 | print(f'Sinkhorn - tfidf: {source_lang} - {target_lang}') | 123 | args = parser.parse_args() |
249 | clf = WassersteinDistances(W_embed=W_common, n_neighbors=5, n_jobs=14, sinkhorn=True) | ||
250 | clf.fit(X_train_idf[:instances], np.ones(instances)) | ||
251 | dist, preds = clf.kneighbors(X_test_idf[:instances], n_neighbors=instances) | ||
252 | mrr, p_at_1 = mrr_precision_at_k(list(range(len(preds))), preds) | ||
253 | print(f'MRR: {mrr} | Precision @ 1: {p_at_1}') | ||
254 | 124 | ||
255 | percentage = p_at_1 * 100 | 125 | main(args) |
256 | fields = [f'{source_lang}', f'{target_lang}', f'{instances}', f'{mrr}', f'{p_at_1}', f'{percentage}'] | ||
257 | with open('/home/syigit/multilang_results/sinkhorn_retrieval_result.csv', 'a') as f: | ||
258 | writer = csv.writer(f) | ||
259 | writer.writerow(fields) | ||