1 files changed, 6 insertions, 9 deletions

diff --git a/WMD_retrieval.py b/WMD_retrieval.py
index 3328023..02f35be 100644
--- a/WMD_retrieval.py
+++ b/WMD_retrieval.py
@@ -6,7 +6,7 @@ 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 (WassersteinRetriever,
                                   clean_corpus_using_embeddings_vocabulary,
                                   load_embeddings)
 
@@ -101,16 +101,13 @@ def main(args):
 
     for metric in runfor:
         if not batch:
-            print(f'{metric} - tfidf: {source_lang} - {target_lang}')
+            print(f'{metric}: {source_lang} - {target_lang}')
 
-        clf = Wasserstein_Retriever(W_embed=W_common,
-                                    n_neighbors=5,
-                                    n_jobs=14,
-                                    sinkhorn=(metric == 'snk'))
+        clf = WassersteinRetriever(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)