Merge branch 'master' of github.com:experimaestro/datamaestro_text

Author: bpiwowar

src/datamaestro_text/data/ir/__init__.py

@@ -1,6 +1,7 @@
 """Generic data types for information retrieval"""
 
 from abc import ABC, abstractmethod
+from enum import Enum
 from functools import cached_property
 import logging
 from pathlib import Path
@@ -88,6 +89,19 @@ def document_recordtype(self) -> Type[DocumentRecord]:
         ...
 
 
+class FileAccess(Enum):
+    """Defines how to access files (e.g. for document stores)"""
+
+    FILE = 0
+    """Direct file access"""
+
+    MMAP = 1
+    """Use mmap"""
+
+    MEMORY = 2
+    """Use memory"""
+
+
 class DocumentStore(Documents):
     """A document store
 
@@ -97,6 +111,10 @@ class DocumentStore(Documents):
     - return the number of documents
     """
 
+    file_access: Meta[FileAccess] = FileAccess.MMAP
+    """How to access the file collection (might not have any impact, depends on
+    the docstore)"""
+
     def docid_internal2external(self, docid: int):
         """Converts an internal collection ID (integer) to an external ID"""
         raise NotImplementedError(f"For class {self.__class__}")
@@ -327,5 +345,4 @@ class PairwiseSampleDataset(Base, ABC):
     """Datasets where each record is a query with positive and negative samples"""
 
     @abstractmethod
-    def iter(self) -> Iterator[PairwiseSample]:
-        ...
+    def iter(self) -> Iterator[PairwiseSample]: ...

src/datamaestro_text/datasets/irds/data.py

@@ -215,7 +215,19 @@ def documentcount(self):
 
     @cached_property
     def store(self):
-        return self.dataset.docs_store()
+        kwargs = {}
+        try:
+            # Translate to ir datasets docstore options
+            import ir_datasets.indices as ir_indices
+            file_access = {
+                ir.FileAccess.MMAP: ir_indices.FileAccess.MMAP,
+                ir.FileAccess.FILE: ir_indices.FileAccess.FILE,
+                ir.FileAccess.MEMORY: ir_indices.FileAccess.MEMORY
+            }[self.file_access]
+            kwargs = {"options": ir_indices.DocstoreOptions(file_access=file_access)}
+        except ImportError:
+            logging.warning("This version of ir-datasets cannot handle docstore options")
+        return self.dataset.docs_store(**kwargs)
 
     @cached_property
     def _docs(self):