data-forager 0.1.6__py3-none-any.whl → 0.2.0__py3-none-any.whl

data_forager/datasets/tokens_with_aux.py

@@ -0,0 +1,91 @@
+"""
+Dataset for reading tokenized samples with auxiliary data.
+
+This module provides TokensWithAuxDataset which reads samples containing
+both tokens and auxiliary data (e.g., loss masks) using the schema stored
+in the sample index.
+"""
+
+from typing import Dict, Optional
+
+import numpy as np
+
+from data_forager.sample_index import SampleIndex
+from data_forager.datasets.common import Dataset
+
+
+class TokensWithAuxDataset(Dataset):
+    """
+    Dataset that returns tokens with auxiliary data.
+
+    Reads samples containing multiple arrays (tokens + auxiliary data) using
+    the schema from the sample index to parse the concatenated bytes.
+
+    Requires sample_schema in the SampleIndex. Use TokensDataset for indexes
+    without auxiliary data.
+    """
+
+    @classmethod
+    def create_from_index_on_filesystem(
+        cls,
+        base_path: str,
+        name: Optional[str] = None,
+    ) -> "TokensWithAuxDataset":
+        """
+        Create a TokensWithAuxDataset from an index stored on the filesystem.
+
+        :param base_path: Base path where the index is stored.
+        :param name: Optional name for logging.
+
+        :return: TokensWithAuxDataset instance.
+        """
+        from data_forager.index_stores.fs_based import IndexStore
+
+        index_store = IndexStore(base_path=base_path)
+        sample_index = index_store.load()
+
+        return cls(sample_index=sample_index, name=name)
+
+    def __init__(
+        self,
+        sample_index: SampleIndex,
+        name: Optional[str] = None,
+        **kwargs,
+    ):
+        """
+        Initialize the dataset.
+
+        :param sample_index: SampleIndex with sample_schema describing the
+            structure of samples.
+        :param name: Optional name for logging.
+
+        :raises ValueError: If sample_index has no sample_schema.
+        """
+        super().__init__(sample_index, name=name, **kwargs)
+
+        if sample_index.sample_schema is None:
+            raise ValueError(
+                "SampleIndex has no sample_schema. "
+                "Use TokensDataset for indexes without auxiliary data."
+            )
+
+        self._schema = sample_index.sample_schema
+
+    def _process_sample(self, sample_bytes: bytes) -> Dict[str, np.ndarray]:
+        """
+        Parse concatenated bytes into named arrays.
+
+        :param sample_bytes: Raw bytes containing all arrays concatenated.
+
+        :return: Dict mapping array names to numpy arrays.
+        """
+        result = {}
+        for array_spec in self._schema.arrays:
+            dtype = np.dtype(array_spec.dtype)
+            start = array_spec.offset
+            length = self._schema.sample_size * dtype.itemsize
+            result[array_spec.name] = np.frombuffer(
+                sample_bytes[start:start + length],
+                dtype=dtype,
+            )
+        return result

data_forager/index_stores/fs_based.py

@@ -1,3 +1,11 @@
+"""
+Filesystem-based index store for persisting sample indexes.
+
+This module provides IndexStore which saves and loads sample indexes
+to/from the filesystem.
+"""
+
+import json
 import os.path
 import shutil
 from typing import Optional, TextIO, BinaryIO
@@ -8,19 +16,28 @@ from basics.base import Base
 
 from data_forager.index_stores.common import IndexStoreInterface
 from data_forager.sample_index import SampleIndex
+from data_forager.sample_generators.schema import SampleSchema, ArraySpec
 
 
 class IndexStore(Base, IndexStoreInterface):
+    """
+    Filesystem-based index store for saving and loading sample indexes.
+
+    Stores index data in a directory with the following structure:
+    - file_location.txt: List of file paths (relative to base_path)
+    - sample_locations.bin: Binary array of (file_index, byte_offset, num_bytes)
+    - sample_schema.json: Optional schema for structured samples
+    """
 
     def __init__(self, base_path: str, index_data_folder: str = "index", name: Optional[str] = None):
         """
+        Initialize the index store.
 
         :param base_path: Base path where the index files are stored.
             File paths in file_location.txt are stored relative to this path.
-
-        :param name: Name of instance, if not provided, the classname will be used
+        :param index_data_folder: Name of the folder within base_path for index files.
+        :param name: Name of instance, if not provided, the classname will be used.
         """
-
         super().__init__(pybase_logger_name=name)
 
         self._base_path = os.path.abspath(base_path)
@@ -31,6 +48,9 @@ class IndexStore(Base, IndexStoreInterface):
         self._file_location_handle: Optional[TextIO] = None
         self._sample_locations_handle: Optional[BinaryIO] = None
 
+        # Optional schema for structured samples
+        self._sample_schema: Optional[SampleSchema] = None
+
     def init_store(self):
         if os.path.exists(self._index_data_path):
             raise ValueError(f"Provided index path already exists: {self._index_data_path}")
@@ -64,6 +84,15 @@ class IndexStore(Base, IndexStoreInterface):
         sample_location_bytes = np.array([file_index, byte_offset, num_bytes], dtype=np.uint64).tobytes()
         self._sample_locations_handle.write(sample_location_bytes)
 
+    def set_sample_schema(self, schema: SampleSchema) -> None:
+        """
+        Set and persist the sample schema.
+
+        :param schema: SampleSchema describing the structure of samples.
+        """
+        self._sample_schema = schema
+        self._save_schema()
+
     def close(self):
         """Close file handles and flush buffered data."""
         if self._file_location_handle is not None:
@@ -78,6 +107,11 @@ class IndexStore(Base, IndexStoreInterface):
         self.close()
 
     def load(self) -> SampleIndex:
+        """
+        Load the sample index from disk.
+
+        :return: SampleIndex with file locations, sample locations, and optional schema.
+        """
         with open(os.path.join(self._index_data_path, "file_location.txt"), "r") as f:
             relative_locations = [loc[:-1] if loc[-1]=='\n' else loc for loc in f.readlines()]
 
@@ -91,7 +125,46 @@ class IndexStore(Base, IndexStoreInterface):
             sample_locations = np.frombuffer(data, dtype=np.uint64)
             sample_locations = sample_locations.reshape((-1, 3))
 
-        return SampleIndex(file_locations, sample_locations)
+        # Load schema if it exists
+        sample_schema = self._load_schema()
+
+        return SampleIndex(file_locations, sample_locations, sample_schema)
+
+    def _save_schema(self) -> None:
+        """Save the sample schema to sample_schema.json."""
+        if self._sample_schema is None:
+            return
+
+        schema_path = os.path.join(self._index_data_path, "sample_schema.json")
+        schema_dict = {
+            "sample_size": self._sample_schema.sample_size,
+            "arrays": [
+                {"name": arr.name, "dtype": arr.dtype, "offset": arr.offset}
+                for arr in self._sample_schema.arrays
+            ],
+            "total_bytes_per_sample": self._sample_schema.total_bytes_per_sample,
+        }
+        with open(schema_path, "w") as f:
+            json.dump(schema_dict, f, indent=2)
+
+    def _load_schema(self) -> Optional[SampleSchema]:
+        """Load the sample schema from sample_schema.json if it exists."""
+        schema_path = os.path.join(self._index_data_path, "sample_schema.json")
+        if not os.path.exists(schema_path):
+            return None
+
+        with open(schema_path, "r") as f:
+            schema_dict = json.load(f)
+
+        arrays = [
+            ArraySpec(name=arr["name"], dtype=arr["dtype"], offset=arr["offset"])
+            for arr in schema_dict["arrays"]
+        ]
+        return SampleSchema(
+            sample_size=schema_dict["sample_size"],
+            arrays=arrays,
+            total_bytes_per_sample=schema_dict["total_bytes_per_sample"],
+        )
 
     def exists(self) -> bool:
         """Check if the index already exists."""

data_forager/indexers/text_lines.py

@@ -1,5 +1,12 @@
-from typing import Optional, List, Tuple, Protocol
-from dataclasses import dataclass
+"""
+File text lines indexer for building sample indexes.
+
+This module provides FileTextLinesIndexer which scans text files line by line,
+uses a sample generator to transform lines into samples, and builds an index
+for random access.
+"""
+
+from typing import Optional, List
 
 import os
 import sys
@@ -8,77 +15,19 @@ from basics.base import Base
 from tqdm import tqdm
 
 from data_forager.index_stores.common import IndexStoreInterface
-
-
-@dataclass
-class SampleData:
-
-    sample_bytes: bytes
-    file_path: str
-
-class SampleGeneratorInterface(Protocol):
-
-    def prepare(self, text_file_path: str):
-        """
-        Prepare sample generation from a new input text file
-
-        :param text_file_path: path to text file
-
-        :return:
-        """
-        ...
-
-    def create_samples(self, text_line: bytes) -> List[SampleData]:
-        """
-        Creates one or more samples from the given text_line and stores it in one or multiple different files.
-        The path to the file(s) in which the samples are stores are also returned.
-
-        IMPORTANT: it is assumed that each sample returned is stored in a file sequentially in the same order.
-                   This must also hold over multiple function calls. This is important because the byte offset
-                   of a sample is derived from the order the samples are returned.
-
-        :param text_line: Text line in bytes from text_file_path, provided in the prepare phase.
-                          The function needs to choose a text encoding itself
-
-        :return: List of DataSample objects. For each created sample the following is given:
-            * Its representation in bytes, as used to store the sample
-            * The file path to where the sample is stored
-
-        """
-        ...
-
-    def finish(self, is_last_file: bool):
-        """
-        Finish generation of samples from text lines of input file at the `text_file_path` given in the prepare() phase.
-
-        is_last_file: indicates if the input text file was the last file to be processed
-
-        :return:
-        """
-        ...
-
-
-class NOOPSampleGenerator(SampleGeneratorInterface):
-
-    def __init__(self):
-        self._current_text_file = None
-
-    def prepare(self, text_file_path: str):
-        self._current_text_file = text_file_path
-
-    def create_samples(self, text_line: bytes) -> List[SampleData]:
-        return [SampleData(text_line, self._current_text_file)]
-
-    def finish(self, is_last_file: bool):
-        self._current_text_file = None
-
-
-def noop_sample_processing(text_line: bytes, text_file_path: str) -> List[SampleData]:
-
-    return [SampleData(text_line, text_file_path)]
+from data_forager.sample_generators.common import (
+    SampleGeneratorInterface,
+    NOOPSampleGenerator,
+)
 
 
 class FileTextLinesIndexer(Base):
+    """
+    Indexes text files by scanning lines and building a sample index.
+
+    Uses a sample generator to transform text lines into samples, then records
+    byte offsets in an index store for O(1) random access during training.
+    """
 
     def __init__(
         self,
@@ -104,13 +53,19 @@ class FileTextLinesIndexer(Base):
 
     def __call__(self):
         """
-        IMPORTANT: input files are always read in binary mode; applying a text encoding is up to the user.
-                   E.g. through process_sample_func and/or when processing the data Dataset::_process_sample()
+        Run the indexing process.
 
-        :return:
+        IMPORTANT: Input files are always read in binary mode; applying a text
+        encoding is up to the user, e.g., through process_sample_func and/or
+        when processing the data in Dataset._process_sample().
         """
         self._index_store.init_store()
 
+        # Set schema if the sample generator provides one
+        sample_schema = self._sample_generator.get_sample_schema()
+        if sample_schema is not None:
+            self._index_store.set_sample_schema(sample_schema)
+
         byte_offset_map = {}
 
         for input_file_path in self._input_file_paths: