1 files changed, 73 insertions, 39 deletions

diff --git a/WMD_retrieval.py b/WMD_retrieval.py
index f32372f..3328023 100644
--- a/WMD_retrieval.py
+++ b/WMD_retrieval.py
@@ -1,15 +1,19 @@
 import argparse
-import numpy as np
+import csv
 import random
+
+import numpy as np
 from sklearn.feature_extraction.text import CountVectorizer, TfidfVectorizer
 from sklearn.preprocessing import normalize
-from Wasserstein_Distance import Wasserstein_Retriever
-from Wasserstein_Distance import load_embeddings, clean_corpus_using_embeddings_vocabulary, mrr_precision_at_k
-import csv
+
+from Wasserstein_Distance import (Wasserstein_Retriever,
+                                  clean_corpus_using_embeddings_vocabulary,
+                                  load_embeddings)
+
 
 def main(args):
 
-    np.seterr(divide='ignore') # POT has issues with divide by zero errors
+    np.seterr(divide='ignore')  # POT has issues with divide by zero errors
     source_lang = args.source_lang
     target_lang = args.target_lang
 
@@ -25,32 +29,38 @@ def main(args):
     mode = args.mode
     runfor = list()
 
-    if (mode == 'all'):
-        runfor.extend(['wmd','snk'])
+    if mode == 'all':
+        runfor.extend(['wmd', 'snk'])
     else:
         runfor.append(mode)
 
-    defs_source = [line.rstrip('\n') for line in open(source_defs_filename, encoding='utf8')]
-    defs_target = [line.rstrip('\n') for line in open(target_defs_filename, encoding='utf8')]
+    defs_source = [
+        line.rstrip('\n')
+        for line in open(source_defs_filename, encoding='utf8')
+    ]
+    defs_target = [
+        line.rstrip('\n')
+        for line in open(target_defs_filename, encoding='utf8')
+    ]
 
     clean_src_corpus, clean_src_vectors, src_keys = clean_corpus_using_embeddings_vocabulary(
-            set(vectors_source.keys()),
-            defs_source,
-            vectors_source,
-            source_lang,
-            )
+        set(vectors_source.keys()),
+        defs_source,
+        vectors_source,
+        source_lang,
+    )
 
     clean_target_corpus, clean_target_vectors, target_keys = clean_corpus_using_embeddings_vocabulary(
-            set(vectors_target.keys()),
-            defs_target,
-            vectors_target,
-            target_lang,
-            )
+        set(vectors_target.keys()),
+        defs_target,
+        vectors_target,
+        target_lang,
+    )
 
     take = args.instances
 
     common_keys = set(src_keys).intersection(set(target_keys))
-    take = min(len(common_keys), take) # you can't sample more than length
+    take = min(len(common_keys), take)  # you can't sample more than length
     experiment_keys = random.sample(common_keys, take)
 
     instances = len(experiment_keys)
@@ -58,13 +68,18 @@ def main(args):
     clean_src_corpus = list(clean_src_corpus[experiment_keys])
     clean_target_corpus = list(clean_target_corpus[experiment_keys])
 
-    if (not batch):
-        print(f'{source_lang} - {target_lang} : document sizes: {len(clean_src_corpus)}, {len(clean_target_corpus)}')
+    if not batch:
+        print(
+            f'{source_lang} - {target_lang} : document sizes: {len(clean_src_corpus)}, {len(clean_target_corpus)}'
+        )
 
     del vectors_source, vectors_target, defs_source, defs_target
 
     vec = CountVectorizer().fit(clean_src_corpus + clean_target_corpus)
-    common = [word for word in vec.get_feature_names() if word in clean_src_vectors or word in clean_target_vectors]
+    common = [
+        word for word in vec.get_feature_names()
+        if word in clean_src_vectors or word in clean_target_vectors
+    ]
     W_common = []
     for w in common:
         if w in clean_src_vectors:
@@ -72,8 +87,10 @@ def main(args):
         else:
             W_common.append(np.array(clean_target_vectors[w]))
 
-    if (not batch):
-        print(f'{source_lang} - {target_lang}: the vocabulary size is {len(W_common)}')
+    if not batch:
+        print(
+            f'{source_lang} - {target_lang}: the vocabulary size is {len(W_common)}'
+        )
 
     W_common = np.array(W_common)
     W_common = normalize(W_common)
@@ -82,26 +99,28 @@ def main(args):
     X_train_idf = vect.transform(clean_src_corpus)
     X_test_idf = vect.transform(clean_target_corpus)
 
-    vect_tf = CountVectorizer(vocabulary=common, dtype=np.double)
-    vect_tf.fit(clean_src_corpus + clean_target_corpus)
-    X_train_tf = vect_tf.transform(clean_src_corpus)
-    X_test_tf = vect_tf.transform(clean_target_corpus)
-
     for metric in runfor:
-        if (not batch):
+        if not batch:
             print(f'{metric} - tfidf: {source_lang} - {target_lang}')
 
-        clf = Wasserstein_Retriever(W_embed=W_common, n_neighbors=5, n_jobs=14, sinkhorn=(metric == 'snk'))
+        clf = Wasserstein_Retriever(W_embed=W_common,
+                                    n_neighbors=5,
+                                    n_jobs=14,
+                                    sinkhorn=(metric == 'snk'))
         clf.fit(X_train_idf[:instances], np.ones(instances))
         # dist, preds = clf.kneighbors(X_test_idf[:instances], n_neighbors=instances)
         # mrr, p_at_one = mrr_precision_at_k(list(range(len(preds))), preds)
         # percentage = p_at_one * 100
-        p_at_one, percentage  = clf.align(X_test_idf[:instances], n_neighbors=instances)
+        p_at_one, percentage = clf.align(X_test_idf[:instances],
+                                         n_neighbors=instances)
 
-        if (not batch):
+        if not batch:
             print(f'P @ 1: {p_at_one}\ninstances: {instances}\n{percentage}%')
         else:
-            fields = [f'{source_lang}', f'{target_lang}', f'{instances}', f'{p_at_one}', f'{percentage}']
+            fields = [
+                f'{source_lang}', f'{target_lang}', f'{instances}',
+                f'{p_at_one}', f'{percentage}'
+            ]
             with open(f'{metric}_retrieval_result.csv', 'a') as f:
                 writer = csv.writer(f)
                 writer.writerow(fields)
@@ -109,16 +128,31 @@ def main(args):
 
 if __name__ == "__main__":
 
-    parser = argparse.ArgumentParser(description='run retrieval using wmd or snk')
+    parser = argparse.ArgumentParser(
+        description='run retrieval using wmd or snk')
     parser.add_argument('source_lang', help='source language short name')
     parser.add_argument('target_lang', help='target language short name')
     parser.add_argument('source_vector', help='path of the source vector')
     parser.add_argument('target_vector', help='path of the target vector')
     parser.add_argument('source_defs', help='path of the source definitions')
     parser.add_argument('target_defs', help='path of the target definitions')
-    parser.add_argument('-b', '--batch', action='store_true', help='running in batch (store results in csv) or running a single instance (output the results)')
-    parser.add_argument('mode', choices=['all', 'wmd', 'snk'], default='all', help='which methods to run')
-    parser.add_argument('-n', '--instances', help='number of instances in each language to retrieve', default=1000, type=int)
+    parser.add_argument(
+        '-b',
+        '--batch',
+        action='store_true',
+        help=
+        'running in batch (store results in csv) or running a single instance (output the results)'
+    )
+    parser.add_argument('mode',
+                        choices=['all', 'wmd', 'snk'],
+                        default='all',
+                        help='which methods to run')
+    parser.add_argument(
+        '-n',
+        '--instances',
+        help='number of instances in each language to retrieve',
+        default=1000,
+        type=int)
 
     args = parser.parse_args()