Merge pull request #5 from MathVast/master

Add DprW100 converters, MsMarcoV2Passage and Trec-News

Author: bpiwowar

src/datamaestro_text/data/ir/base.py

@@ -43,6 +43,12 @@ class IDItem(Item, ABC):
 
     id: str
 
+@define
+class UrlItem(Item):
+    """An url item"""
+
+    url: str
+
 
 @define
 class AdhocAssessment:

src/datamaestro_text/data/ir/formats.py

@@ -58,7 +58,7 @@ class NFCorpusDocument(TextItem):
 
     @cached_property
     def text(self):
-        return self.abstract
+        return f"{self.title} {self.abstract}"
 
 
 @define
@@ -99,8 +99,8 @@ class WapoDocument(TextItem):
     body_media: Tuple[WapoDocMedia, ...]
 
     @cached_property
-    def text(self):
-        return self.body
+    def text(self):        
+        return f"{self.title} {self.body_paras_html}"
 
 
 @define
@@ -128,6 +128,15 @@ def text(self):
 
 
 @define
+class DprW100Doc(TextItem):
+    text: str
+    title: str
+
+@define 
+class MsMarcoV2Passage(TextItem):
+    text: str
+    spans: Tuple[Tuple[int, int], ...]
+    msmarco_document_id: str
 class Touche2020(TextItem):
     text: str
     title: str
@@ -194,3 +203,17 @@ class TrecTopic(SimpleTextItem):
 
 
 TrecTopicRecord = record_type(IDItem, TrecTopic)
+
+@define
+class DprW100Query(TextItem):
+    text: str
+    answers: Tuple[str]
+
+@define
+class TrecBackgroundLinkingQuery(IDItem):
+    query_id: str
+    doc_id: str
+    url: str
+
+    def get_text(self):
+        raise NotImplementedError()

src/datamaestro_text/datasets/irds/data.py

@@ -37,6 +37,7 @@
     SimpleAdhocAssessment,
     SimpleTextItem,
     TopicRecord,
+    UrlItem,
     create_record,
 )
 
@@ -165,6 +166,19 @@ class Documents(ir.DocumentStore, IRDSId):
             "source",
             "source_content_type",
         ),
+        _irds.dpr_w100.DprW100Doc: tuple_constructor(
+            formats.DprW100Doc,
+            "doc_id",
+            "text",
+            "title",
+        ),
+        _irds.msmarco_passage_v2.MsMarcoV2Passage: tuple_constructor(
+            formats.MsMarcoV2Passage,
+            "doc_id",
+            "text",
+            "spans",
+            "msmarco_document_id",
+        ),
     }
 
     """Wraps an ir datasets collection -- and provide a default text
@@ -385,6 +399,12 @@ class Topics(ir.TopicsStore, IRDSId):
             "tweet_time",
             "description",
         ),
+        _irds.dpr_w100.DprW100Query: tuple_constructor(
+            formats.DprW100Query,
+            "query_id",
+            "text",
+            "answers"
+        ),
     }
 
     HANDLERS = {
@@ -415,7 +435,52 @@ def topic_ext(self, external_topic_id: str) -> TopicRecord:
     def iter(self) -> Iterator[TopicRecord]:
         """Returns an iterator over topics"""
         return self.handler.iter()
+    
+class TrecBackgroundLinkingTopicsHandler(TopicsHandler):
+    def __init__(self, dataset):
+        self.dataset = dataset
+    
+    @cached_property
+    def ext2records(self):
+        return {record[IDItem].id: record for record in self.records}
+
+    def topic_int(self, internal_topic_id: int) -> TopicRecord:
+        """Returns a document given its internal ID"""
+        return self.records[internal_topic_id]
+
+    def topic_ext(self, external_topic_id: str) -> TopicRecord:
+        """Returns a document given its external ID"""
+        return self.ext2records[external_topic_id]
 
+    def iter(self) -> Iterator[ir.TopicRecord]:
+        """Returns an iterator over topics"""
+        return iter(self.records)
+
+    @cached_property
+    def records(self):
+        try:
+            records = []
+
+            for query in self.dataset.dataset.queries_iter():
+                topic =  Record(
+                    IDItem(query.query_id),
+                    # Following BEIR documentation, we use title of documents as queries: https://github.com/beir-cellar/beir/blob/main/examples/dataset/README.md#queries-and-qrels
+                    SimpleTextItem(self.dataset.dataset.docs_store().get(query.doc_id).title), 
+                    UrlItem(query.url),
+                )
+                records.append(topic)
+        except Exception:
+            logging.exception("Error while computing topic records")
+            raise
+
+        return records
+
+
+Topics.HANDLERS.update(
+    {
+        _irds.wapo.TrecBackgroundLinkingQuery: TrecBackgroundLinkingTopicsHandler
+    }
+)
 
 class CastTopicsHandler(TopicsHandler):
     def __init__(self, dataset):