add trec writer (#85)

castorini · Sep 12, 2020 · a258c13 · a258c13
1 parent 4019c3f
commit a258c13
Showing 1 changed file with 10 additions and 1 deletion.
diff --git a/pygaggle/model/writer.py b/pygaggle/model/writer.py
@@ -8,11 +8,12 @@
 
 
 class Writer:
-    def __init__(self, path: Optional[Path] = None, overwrite: bool = True):
+    def __init__(self, path: Optional[Path] = None, overwrite: bool = True, tag: Optional[str] = None):
         self.to_output = str(path) not in [".", None]
         print(f'Writing run: {self.to_output}')
         if self.to_output:
             self.f = open(path, "w" if overwrite else "w+")
+        self.tag = tag
 
     def write_line(self, text: str):
         if self.to_output:
@@ -29,3 +30,11 @@ def write(self, scores: List[float], example: RelevanceExample):
                             key=lambda x: x[1], reverse=True)
         for ct, (doc, score) in enumerate(doc_scores):
             self.write_line(f"{example.query.id}\t{doc.metadata['docid']}\t{ct+1}")
+
+
+class TrecWriter(Writer):
+    def write(self, scores: List[float], example: RelevanceExample):
+        doc_scores = sorted(list(zip(example.documents, scores)),
+                            key=lambda x: x[1], reverse=True)
+        for ct, (doc, score) in enumerate(doc_scores):
+            self.write_line(f"{example.query.id}\tQ0\t{doc.metadata['docid']}\t{ct+1}\t{score}\t{self.tag}")