Add saving results to file

.gitignore

@@ -7,6 +7,7 @@
 
 # Data
 dataset/
+results/
 *.p
 *.png
 

retrieval/util/FileManager.py

@@ -1,4 +1,5 @@
 import csv
+import os
 from typing import Union, Any
 import pickle
 
@@ -10,6 +11,15 @@ def read_tsv(file_name: str) -> Union[list, str]:
     raise RuntimeError("Invalid file type: only '.tsv' files are accepted.")
 
 
+def write_txt(file_name: str, data: str, mode='w'):
+    directory = os.path.dirname(file_name)
+    if not os.path.exists(directory):
+        os.makedirs(directory)
+    file = open(file_name, mode=mode)
+    file.write(data)
+    file.close()
+
+
 def write_pickle(data: object, file_name: str):
     pickle.dump(data, open(file_name, "wb"))
 

start.py

@@ -1,6 +1,7 @@
 import argparse
 
 from retrieval.DatasetParser import DatasetParser
+from retrieval.util.FileManager import write_txt
 
 
 def main():
@@ -11,8 +12,24 @@ def main():
     parser.add_argument('-s', '--smoothing', help='smoothing for the Query Likelihood model')
 
     args = parser.parse_args()
-    parser = DatasetParser(args.dataset)
-    results = parser.parse(args.model, plot_freq=args.plot, smoothing=args.smoothing)
+    dataset = args.dataset
+    model = args.model
+    smoothing = args.smoothing
+    plot = args.plot
+    parser = DatasetParser(dataset)
+    results = parser.parse(model, plot_freq=plot, smoothing=smoothing)
+
+    models = {'bm25': 'BM25', 'vector': 'VS', 'query': 'LM'}
+    smoothers = {'laplace': '-Laplace', 'lidstone': '-Lidstone', 'dirichlet': '-Dirichlet'}
+    model = models[model]
+    smoothing = smoothers[smoothing] if smoothing is not None else ""
+
+    data = ''
+    for qid, passages in results.items():
+        for rank, (pid, score) in enumerate(passages.items()):
+            data += f"{qid}\t{'A1'}\t{pid}\t{rank}\t{format(score, '.2f')}\t{model}{smoothing}\n"
+
+    write_txt(f'results/{model}.txt', data)
 
 
 if __name__ == '__main__':