Cleaned up TextItem hierarchy

experimaestro · Feb 27, 2024 · 6d73bcf · 6d73bcf
1 parent 9240d8c
commit 6d73bcf
Show file tree

Hide file tree

Showing 5 changed files with 48 additions and 86 deletions.
diff --git a/src/datamaestro_text/data/ir/base.py b/src/datamaestro_text/data/ir/base.py
@@ -6,12 +6,12 @@
 
 class BaseRecord(Record):
     @classmethod
-    def from_text(cls, text: str):
-        return cls(SimpleTextItem(text))
+    def from_text(cls, text: str, *items: Item):
+        return cls(SimpleTextItem(text), *items)
 
     @classmethod
-    def from_id(cls, id: str):
-        return cls(IDItem(id))
+    def from_id(cls, id: str, *items: Item):
+        return cls(IDItem(id), *items)
 
 
 class TopicRecord(BaseRecord):
@@ -32,20 +32,18 @@ class ScoredItem(Item):
 
 
 class TextItem(Item, ABC):
+    @property
     @abstractmethod
-    def get_text(self) -> str:
+    def text(self) -> str:
         """Returns the text"""
 
 
 @define
-class SimpleTextItem(TextItem, ABC):
+class SimpleTextItem(TextItem):
     """A topic/document with a text record"""
 
     text: str
 
-    def get_text(self):
-        return self.text
-
 
 @define
 class InternalIDItem(Item, ABC):

diff --git a/src/datamaestro_text/data/ir/cord19.py b/src/datamaestro_text/data/ir/cord19.py
@@ -48,5 +48,5 @@ def iter(self) -> Iterator[CordDocumentRecord]:
             for row in DictReader(fp):
                 yield CordDocumentRecord(
                     IDItem(row["cord_uid"]),
-                    DocumentWithTitle(row["title"], row["abstract"]),
+                    DocumentWithTitle(row["abstract"], row["title"]),
                 )
diff --git a/src/datamaestro_text/data/ir/formats.py b/src/datamaestro_text/data/ir/formats.py
@@ -1,3 +1,4 @@
+from functools import cached_property
 from typing import ClassVar, Tuple
 from attrs import define
 from datamaestro.record import recordtypes
@@ -7,25 +8,22 @@
 
 
 @define
-class CordDocument(TextItem):
-    text: str
-    title: str
-    url: str
-    pubmed_id: str
-
-    has_text: ClassVar[bool] = True
+class DocumentWithTitle(TextItem):
+    """Web document with title and body"""
 
-    def get_text(self):
-        return f"{self.title} {self.text}"
+    body: str
 
+    title: str
 
-@define
-class DocumentWithTitle(TextItem):
-    """Web document with title and URL"""
+    @cached_property
+    def text(self):
+        return f"{self.title} {self.body}"
 
-    title: str
 
-    text: str
+@define
+class CordDocument(DocumentWithTitle):
+    url: str
+    pubmed_id: str
 
 
 @define
@@ -36,10 +34,9 @@ class CordFullTextDocument(TextItem):
     abstract: str
     body: Tuple[Cord19FullTextSection, ...]
 
-    has_text: ClassVar[bool] = True
-
-    def get_text(self):
-        return f"{self.abstract}"
+    @cached_property
+    def text(self):
+        return self.abstract
 
 
 @define
@@ -48,10 +45,9 @@ class MsMarcoDocument(TextItem):
     title: str
     body: str
 
-    has_text: ClassVar[bool] = True
-
-    def get_text(self):
-        return f"{self.body}"
+    @cached_property
+    def text(self):
+        return self.body
 
 
 @define
@@ -60,31 +56,24 @@ class NFCorpusDocument(TextItem):
     title: str
     abstract: str
 
-    has_text: ClassVar[bool] = True
-
-    def get_text(self):
-        return f"{self.abstract}"
+    @cached_property
+    def text(self):
+        return self.abstract
 
 
 @define
 class TitleDocument(TextItem):
-    text: str
+    body: str
     title: str
-    has_text: ClassVar[bool] = True
 
-    def get_text(self):
-        return f"{self.title} {self.text}"
+    @cached_property
+    def text(self):
+        return f"{self.title} {self.body}"
 
 
 @define
-class TitleUrlDocument(TextItem):
-    text: str
-    title: str
+class TitleUrlDocument(TitleDocument):
     url: str
-    has_text: ClassVar[bool] = True
-
-    def get_text(self):
-        return f"{self.title} {self.text}"
 
 
 @define
@@ -93,9 +82,8 @@ class TrecParsedDocument(TextItem):
     body: str
     marked_up_doc: bytes
 
-    has_text: ClassVar[bool] = True
-
-    def get_text(self):
+    @cached_property
+    def text(self):
         return f"{self.title} {self.body}"
 
 
@@ -110,10 +98,9 @@ class WapoDocument(TextItem):
     body_paras_html: Tuple[str, ...]
     body_media: Tuple[WapoDocMedia, ...]
 
-    has_text: ClassVar[bool] = True
-
-    def get_text(self):
-        return f"{self.body}"
+    @cached_property
+    def text(self):
+        return self.body
 
 
 @define
@@ -127,22 +114,18 @@ class TweetDoc(TextItem):
     source: bytes
     source_content_type: str
 
-    def get_text(self):
-        return f"{self.text}"
-
 
 @define
 class OrConvQADocument(TextItem):
     id: str
     title: str
-    text: str
+    body: str
     aid: str
     bid: int
 
-    has_text: ClassVar[bool] = True
-
-    def get_text(self):
-        return f"{self.title} {self.text}"
+    @cached_property
+    def text(self):
+        return f"{self.title} {self.body}"
 
 
 @define
@@ -151,37 +134,18 @@ class TrecTopic(TextItem):
     query: str
     narrative: str
 
-    def get_text(self):
-        return f"{self.text}"
-
 
 @define
 class UrlTopic(TextItem):
     text: str
     url: str
 
-    def get_text(self):
-        return f"{self.text}"
-
 
 @define
 class NFCorpusTopic(TextItem):
-    title: str
+    text: str
     all: str
 
-    def get_text(self):
-        return f"{self.title}"
-
-
-@define
-class TrecQuery(TextItem):
-    title: str
-    description: str
-    narrative: str
-
-    def get_text(self):
-        return f"{self.description}"
-
 
 @define
 class TrecMb13Query(TextItem):

diff --git a/src/datamaestro_text/datasets/irds/data.py b/src/datamaestro_text/datasets/irds/data.py
@@ -212,7 +212,7 @@ def converter(self):
 
 if hasattr(_irds, "miracl"):
     Documents.CONVERTERS[_irds.miracl.MiraclDoc] = tuple_constructor(
-        formats.DocumentWithTitle, "doc_id", "title", "text"
+        formats.DocumentWithTitle, "doc_id", "text", "title"
     )
 
 
@@ -351,7 +351,7 @@ class Topics(ir.TopicsStore, IRDSId):
             formats.NFCorpusTopic, "query_id", "title", "all"
         ),
         TrecQuery: tuple_constructor(
-            formats.TrecQuery, "query_id", "title", "description", "narrative"
+            formats.TrecTopic, "query_id", "title", "description", "narrative"
         ),
         _irds.tweets2013_ia.TrecMb13Query: tuple_constructor(
             formats.TrecMb13Query, "query_id", "query", "time", "tweet_time"

diff --git a/src/datamaestro_text/transforms/ir/__init__.py b/src/datamaestro_text/transforms/ir/__init__.py
@@ -170,7 +170,7 @@ def get_query(query):
         else:
 
             def get_query(query):
-                return query[ir.TextItem].get_text()
+                return query[ir.TextItem].text
 
         if self.doc_ids:
 
@@ -180,7 +180,7 @@ def get_doc(doc):
         else:
 
             def get_doc(doc):
-                return doc.get_text()
+                return doc.text
 
         def triplegenerator():
             logging.info("Starting to output triples")