feat: TREC adhoc runs read/write

Author: bpiwowar

src/datamaestro_text/data/ir/__init__.py

@@ -5,7 +5,7 @@
 import logging
 from pathlib import Path
 from attrs import define
-from typing import Callable, Dict, Iterator, List, Optional, Tuple, Type
+from typing import Callable, Dict, Iterator, List, Optional, Tuple, Type, TYPE_CHECKING
 import random
 from experimaestro import Config
 from datamaestro.definitions import datatasks, Param, Meta
@@ -29,6 +29,9 @@
     AdhocAssessedTopic,
 )
 
+#: A adhoc run dictionary (query id -> doc id -> score)
+AdhocRunDict = dict[str, dict[str, float]]
+
 
 class Documents(Base):
     """A set of documents with identifiers
@@ -185,7 +188,10 @@ def iter(self) -> Iterator[AdhocAssessedTopic]:
 class AdhocRun(Base):
     """IR adhoc run"""
 
-    pass
+    @abstractmethod
+    def get_dict(self) -> "AdhocRunDict":
+        """Get the run as a dictionary query ID -> doc ID -> score"""
+        ...
 
 
 class AdhocResults(Base):

src/datamaestro_text/data/ir/base.py

@@ -71,7 +71,7 @@ class AdhocAssessedTopic:
     """List of assessments for this topic"""
 
 
-def create_record(*items: Item, id: str = None, text: str = None):
+def create_record(*items: Item, id: str = None, text: str = None) -> Record:
     """Easy creation of a text/id item"""
     extra_items = []
     if id is not None:

src/datamaestro_text/data/ir/trec.py

@@ -1,9 +1,9 @@
+import re
 from typing import Dict, List, Optional
-from datamaestro.data import Base
 from experimaestro import documentation, Param, Meta
 from pathlib import Path
-from datamaestro.record import Record
 from datamaestro_text.data.ir import (
+    AdhocRunDict,
     Documents,
     Topics,
     AdhocAssessments,
@@ -47,6 +47,11 @@ def iter(self):
 class TrecAdhocRun(AdhocRun):
     path: Param[Path]
 
+    def get_dict(self) -> AdhocRunDict:
+        import datamaestro_text.interfaces.trec as trec
+
+        return trec.parse_run(self.path)
+
 
 class TrecAdhocResults(AdhocResults):
     """Adhoc results (TREC format)"""
@@ -62,8 +67,6 @@ class TrecAdhocResults(AdhocResults):
 
     def get_results(self) -> Dict[str, float]:
         """Returns the results as a dictionary {metric_name: value}"""
-        import re
-
         re_spaces = re.compile(r"\s+")
 
         results = {}

src/datamaestro_text/interfaces/trec.py

@@ -1,7 +1,7 @@
-from attrs import define
 from pathlib import Path
 from typing import Iterator, Optional
 import re
+from datamaestro_text.data.ir import AdhocRunDict
 from datamaestro_text.data.ir.base import (
     AdhocAssessedTopic,
     TopicRecord,
@@ -10,6 +10,33 @@
 )
 from datamaestro_text.data.ir.formats import TrecTopicRecord, TrecTopic
 
+# --- Runs
+
+
+def parse_run(path: Path) -> AdhocRunDict:
+    results = {}
+    with path.open("rt") as f:
+        for line in f:
+            query_id, _q0, doc_id, _rank, score, _model_id = re.split(
+                r"\s+", line.strip()
+            )
+            results.setdefault(query_id, {})[doc_id] = score
+
+    return results
+
+
+def write_run_dict(run: AdhocRunDict, run_path: Path):
+    """Write run dict"""
+    with run_path.open("wt") as f:
+        for query_id, scored_documents in run.items():
+            scored_documents = list(
+                [(doc_id, score) for doc_id, score in scored_documents.items()]
+            )
+            scored_documents.sort(key=lambda x: x[1], reverse=True)
+            for ix, (doc_id, score) in enumerate(scored_documents):
+                f.write(f"{query_id} Q0 {doc_id} {ix + 1} {score} run\n")
+
+
 # --- Assessments
 
 
@@ -59,7 +86,7 @@ def parse_query_format(file, xml_prefix=None) -> Iterator[TopicRecord]:
                 num = line[len("<num>") :].replace("Number:", "").strip()
                 reading = None
             elif line.startswith(f"<{xml_prefix}title>"):
-                title = line[len(f"<{xml_prefix}title>") :].strip()
+                title = line[len(f"<{xml_prefix}title>") : ].strip()
                 if title == "":
                     reading = "title"
                 else: