Merge pull request #3 from yzong12138/beir_adding

update the code: adding the structure for adapting the beir dataset

Author: bpiwowar

src/datamaestro_text/data/ir/formats.py

@@ -1,5 +1,5 @@
 from functools import cached_property
-from typing import ClassVar, Tuple
+from typing import ClassVar, Tuple, List
 from attrs import define
 from datamaestro.record import record_type
 from ir_datasets.datasets.wapo import WapoDocMedia
@@ -126,6 +126,22 @@ class OrConvQADocument(TextItem):
     def text(self):
         return f"{self.title} {self.body}"
 
+@define
+class Touche2020(TextItem):
+    text: str
+    title: str
+    stance: str
+    url: str
+
+@define
+class SciDocs(TextItem):
+    text: str
+    title: str
+    authors: List[str]
+    year: int
+    cited_by: List[str]
+    references: List[str]
+
 
 @define
 class UrlTopic(TextItem):
@@ -159,6 +175,13 @@ class TrecMb14Query(TextItem):
     def get_text(self):
         return f"{self.query}"
 
+@define 
+class SciDocsTopic(TextItem):
+    text: str
+    authors: List[str]
+    year: int
+    cited_by: List[str]
+    references: List[str]
 
 @define()
 class TrecTopic(SimpleTextItem):

src/datamaestro_text/datasets/irds/data.py

@@ -110,6 +110,12 @@ class Documents(ir.DocumentStore, IRDSId):
         _irds.beir.BeirTitleUrlDoc: tuple_constructor(
             formats.TitleUrlDocument, "doc_id", "text", "title", "url"
         ),
+        _irds.beir.BeirToucheDoc: tuple_constructor(
+            formats.Touche2020, "doc_id", "text", "title", "stance", "url"
+        ),
+        _irds.beir.BeirSciDoc: tuple_constructor(
+            formats.SciDocs, "doc_id", "text", "title", "authors", "year", "cited_by", "references"
+        ),
         _irds.msmarco_document.MsMarcoDocument: tuple_constructor(
             formats.MsMarcoDocument, "doc_id", "url", "title", "body"
         ),
@@ -355,6 +361,12 @@ class Topics(ir.TopicsStore, IRDSId):
         TrecQuery: tuple_constructor(
             formats.TrecTopic, "query_id", "title", "description", "narrative"
         ),
+        _irds.beir.BeirToucheQuery: tuple_constructor(
+            formats.TrecTopic, "query_id", "text", "description", "narrative"
+        ),
+        _irds.beir.BeirSciQuery: tuple_constructor(
+            formats.SciDocsTopic, "query_id", "text", "authors", "year", "cited_by", "references"
+        ),
         _irds.tweets2013_ia.TrecMb13Query: tuple_constructor(
             formats.TrecMb13Query, "query_id", "query", "time", "tweet_time"
         ),