From 2936635892e17031c37facfd2115e8cfd6633222 Mon Sep 17 00:00:00 2001 From: Yigit Sever Date: Sun, 22 Sep 2019 01:33:24 +0300 Subject: Introduce linter, stylize --- WMD_retrieval.py | 112 ++++++++++++++++++++++++++++++++++++------------------- 1 file changed, 73 insertions(+), 39 deletions(-) (limited to 'WMD_retrieval.py') 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() -- cgit v1.2.3-70-g09d2