add storing prediction/qrel files option for test set (#125)

* add prediction/qrel files dump option

* fix comment
This commit is contained in:
Victor Yang
2018-06-18 20:43:40 -04:00
committed by GitHub
parent 5cc027f03e
commit 921a45e7ca
8 changed files with 33 additions and 17 deletions
+7 -3
View File
@@ -3,7 +3,7 @@ import subprocess
import time
def get_map_mrr(qids, predictions, labels, device=0):
def get_map_mrr(qids, predictions, labels, device=0, keep_results=False):
"""
Get the map and mrr using the trec_eval utility.
qids, predictions, labels should have the same length.
@@ -30,7 +30,11 @@ def get_map_mrr(qids, predictions, labels, device=0):
mean_average_precision = float(trec_out_lines[0].split('\t')[-1])
mean_reciprocal_rank = float(trec_out_lines[1].split('\t')[-1])
os.remove(qrel_fname)
os.remove(results_fname)
if keep_results:
print("Saving prediction file to {}".format(results_fname))
print("Saving qrel file to {}".format(qrel_fname))
else:
os.remove(results_fname)
os.remove(qrel_fname)
return mean_average_precision, mean_reciprocal_rank