Merge pull request #71 from fschlatt/main

Return none for dataset len if no length can be inferred

Author: fschlatt

lightning_ir/data/dataset.py

@@ -199,14 +199,14 @@ def __init__(self, query_dataset: str, num_queries: int | None = None) -> None:
         super(IRDataset, self).__init__()
         self.num_queries = num_queries
 
-    def __len__(self) -> int:
+    def __len__(self) -> int | None:
         """Number of queries in the dataset.
 
         :return: Number of queries
         :rtype: int
         """
         # TODO fix len for multi-gpu and multi-worker inference
-        return self.num_queries or self.ir_dataset.queries_count()
+        return self.num_queries or getattr(self.ir_dataset, "queries_count", lambda: None)() or None
 
     def __iter__(self) -> Iterator[QuerySample]:
         """Iterate over queries in the dataset.
@@ -253,7 +253,7 @@ def __init__(self, doc_dataset: str, num_docs: int | None = None, text_fields: S
         self.num_docs = num_docs
         self.text_fields = text_fields
 
-    def __len__(self) -> int:
+    def __len__(self) -> int | None:
         """Number of documents in the dataset.
 
         :raises ValueError: If no `num_docs` was not provided in the constructor and the number of documents cannot
@@ -262,9 +262,7 @@ def __len__(self) -> int:
         :rtype: int
         """
         # TODO fix len for multi-gpu and multi-worker inference
-        num_docs = self.num_docs or self.ir_dataset.docs_count()
-        if num_docs is None:
-            raise TypeError("Unable to determine number of documents.")
+        num_docs = self.num_docs or getattr(self.ir_dataset, "doc_count", lambda: None)() or None
         return num_docs
 
     def __iter__(self) -> Iterator[DocSample]: