data-forager 0.1.6__tar.gz → 0.2.0__tar.gz

{data_forager-0.1.6 → data_forager-0.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: data-forager
-Version: 0.1.6
+Version: 0.2.0
 Summary: Enabling random access to large datasets on disk for PyTorch training and other use cases
 Author-email: Freddy Snijder <forager@visionscapers.com>
 License-Expression: MIT

data_forager-0.2.0/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-0.1.6 → data_forager-0.2.0}/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-0.1.6 → data_forager-0.2.0}/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:

data_forager-0.2.0/data_forager/indexers/tokenization_indexer.py

@@ -0,0 +1,310 @@
+"""
+Factory function for creating tokenization and indexing pipelines.
+
+This module provides a convenience function for setting up the complete pipeline
+to tokenize JSONL text files and create an index for random access.
+"""
+
+from typing import Callable, Dict, List, Optional
+
+import json
+import logging
+import os
+
+from basics.logging import get_logger
+
+from data_forager.index_stores.common import IndexStoreInterface
+from data_forager.index_stores.fs_based import IndexStore as FSBasedIndexStore
+from data_forager.indexers.text_lines import FileTextLinesIndexer
+from data_forager.sample_generators.tokenization import (
+    TokenizedSampleGenerator,
+    TokenizerFunc,
+    ProcessTextLineFunc,
+)
+from data_forager.sample_generators.aux.common import Part, AuxDataGenerator
+from data_forager.utils import find_files_recursive, natural_sort
+
+
+ProcessPartsFunc = Callable[[bytes], List[Part]]
+
+
+module_logger = get_logger(os.path.basename(__file__))
+
+
+def get_text_from_jsonl(jsonl_bytes: bytes, text_key: str = "text", text_encoding: str = "utf-8") -> str:
+    """
+    Extract text from a JSONL line.
+
+    :param jsonl_bytes: Raw bytes of the JSONL line.
+    :param text_key: Key in the JSON object containing the text.
+    :param text_encoding: Text encoding to use for decoding.
+
+    :return: The extracted text string.
+    """
+    jsonl_text = jsonl_bytes.decode(text_encoding)
+    data = json.loads(jsonl_text)
+    return data[text_key]
+
+
+def create_tokenize_and_index_jsonl_text_func(
+    tokenizer_func: TokenizerFunc,
+    eos_idx: int,
+    input_base_path: Optional[str] = None,
+    input_file_paths: Optional[List[str]] = None,
+    output_base_path: Optional[str] = None,
+    index_store: Optional[IndexStoreInterface] = None,
+    process_text_line_func: Optional[ProcessTextLineFunc] = None,
+    logger: Optional[logging.Logger] = None,
+    name: Optional[str] = None,
+    **sample_generator_kwargs,
+) -> FileTextLinesIndexer:
+    """
+    Create a pipeline to tokenize text from JSONL files and create an index for random access.
+
+    The pipeline:
+    * Tokenizes text from input JSONL objects
+    * Stores the token data in bin files under "tokenized-samples" folder
+    * Stores index data under "index" folder
+
+    Usage:
+    ```python
+    import tiktoken
+
+    enc = tiktoken.get_encoding("gpt2")
+    def tokenize_text(text: str) -> List[int]:
+        return enc.encode_ordinary(text)
+
+    # Option 1: Scan directory for JSONL files, output to same directory
+    indexer = create_tokenize_and_index_jsonl_text_func(
+        tokenizer_func=tokenize_text,
+        eos_idx=enc.eot_token,
+        input_base_path='./data',
+        sample_size=1024,
+    )
+
+    # Option 2: Explicit input files and output path
+    indexer = create_tokenize_and_index_jsonl_text_func(
+        tokenizer_func=tokenize_text,
+        eos_idx=enc.eot_token,
+        input_file_paths=['./data/train.jsonl'],
+        output_base_path='./output',
+        sample_size=1024,
+    )
+
+    # Run tokenization and indexing
+    indexer()
+    ```
+
+    :param tokenizer_func: Function used to tokenize text.
+    :param eos_idx: EOS token index, known by the used Tokenizer.
+    :param input_base_path: Path to directory containing JSONL files (searched recursively).
+        Used as fallback for output if `output_base_path` is not provided.
+    :param input_file_paths: List of file paths to process. If provided, these are used
+        instead of scanning `input_base_path` for JSONL files.
+    :param output_base_path: Base path for output (index and tokenized samples).
+        If not provided, `input_base_path` is used.
+    :param index_store: Index store to use. If provided, this is used instead of
+        creating a new FSBasedIndexStore.
+    :param process_text_line_func: Function used to process text lines.
+        By default, this converts input JSON lines to dicts and returns the "text" field.
+        See function get_text_from_jsonl().
+    :param logger: Logger to use. If not provided, uses module logger.
+    :param name: Name of the indexer, used for logging purposes.
+    :param sample_generator_kwargs: Other kwargs passed to TokenizedSampleGenerator
+        (e.g., sample_size, token_dtype, base_output_path).
+
+    :raises ValueError: If both `input_base_path` and `input_file_paths` are None.
+    :raises ValueError: If `index_store` is None and both `output_base_path` and
+        `input_base_path` are None.
+
+    :return: FileTextLinesIndexer instance that can be called to run tokenization
+        and indexing.
+    """
+    if logger is None:
+        logger = module_logger
+
+    # Validate input source
+    if input_base_path is None and input_file_paths is None:
+        raise ValueError(
+            "Either input_base_path or input_file_paths must be provided"
+        )
+
+    # Determine output base path
+    effective_output_base_path = output_base_path or input_base_path
+
+    # Validate output destination
+    if index_store is None and effective_output_base_path is None:
+        raise ValueError(
+            "Either index_store, output_base_path, or input_base_path must be provided "
+            "to determine where to store the index"
+        )
+
+    logger.info(f"Output base path: {effective_output_base_path}")
+
+    if process_text_line_func is None:
+        process_text_line_func = get_text_from_jsonl
+
+    if index_store is None:
+        index_store = FSBasedIndexStore(
+            base_path=effective_output_base_path,
+        )
+
+    if input_file_paths is None:
+        logger.info(f"Scanning for JSONL files in: {input_base_path}")
+        input_file_paths = find_files_recursive(
+            input_base_path,
+            extension_patterns=['*.jsonl', '*.JSONL']
+        )
+        # Assuming numbered files
+        input_file_paths = natural_sort(input_file_paths)
+        logger.info(f"Found {len(input_file_paths)} JSONL file(s)")
+
+    # Set default base_output_path for tokenized samples if not provided in kwargs
+    if 'base_output_path' not in sample_generator_kwargs:
+        default_base_output_path = os.path.join(
+            effective_output_base_path, "tokenized-samples"
+        )
+        logger.info(f"Tokenized samples output path: {default_base_output_path}")
+        sample_generator_kwargs['base_output_path'] = default_base_output_path
+
+    sample_generator = TokenizedSampleGenerator(
+        process_text_line_func=process_text_line_func,
+        tokenizer_func=tokenizer_func,
+        eos_idx=eos_idx,
+        **sample_generator_kwargs
+    )
+
+    return FileTextLinesIndexer(
+        input_file_paths=input_file_paths,
+        index_store=index_store,
+        sample_generator=sample_generator,
+        description="Tokenizing and indexing",
+        name=name,
+    )
+
+
+def create_tokenize_and_index_with_aux_func(
+    process_parts_func: ProcessPartsFunc,
+    tokenizer_func: TokenizerFunc,
+    eos_idx: int,
+    aux_generators: Dict[str, AuxDataGenerator],
+    input_base_path: Optional[str] = None,
+    input_file_paths: Optional[List[str]] = None,
+    output_base_path: Optional[str] = None,
+    index_store: Optional[IndexStoreInterface] = None,
+    logger: Optional[logging.Logger] = None,
+    name: Optional[str] = None,
+    **sample_generator_kwargs,
+) -> FileTextLinesIndexer:
+    """
+    Create a pipeline to tokenize structured samples with auxiliary data.
+
+    This function creates a pipeline that:
+    * Processes structured input (parts with types) from JSONL files
+    * Tokenizes each part and generates auxiliary data (e.g., loss masks)
+    * Stores concatenated token + aux data in bin files
+    * Creates an index with schema for random access
+
+    Usage:
+    ```python
+    from data_forager.sample_generators.aux import Part, LossMaskGenerator
+
+    def parse_parts(line_bytes: bytes) -> List[Part]:
+        data = json.loads(line_bytes.decode('utf-8'))
+        return [Part(type=p['type'], text=p['text']) for p in data['parts']]
+
+    indexer = create_tokenize_and_index_with_aux_func(
+        process_parts_func=parse_parts,
+        tokenizer_func=tokenizer.encode,
+        eos_idx=tokenizer.eos_token_id,
+        aux_generators={'loss_mask': LossMaskGenerator()},
+        input_base_path='./data',
+        sample_size=4096,
+    )
+
+    indexer()
+    ```
+
+    :param process_parts_func: Function to extract typed parts from input bytes.
+        Takes JSONL bytes and returns List[Part].
+    :param tokenizer_func: Function used to tokenize text.
+    :param eos_idx: EOS token index, known by the used Tokenizer.
+    :param aux_generators: Dict mapping names to AuxDataGenerator instances.
+        Example: {'loss_mask': LossMaskGenerator()}
+    :param input_base_path: Path to directory containing JSONL files.
+    :param input_file_paths: List of file paths to process.
+    :param output_base_path: Base path for output (index and tokenized samples).
+    :param index_store: Index store to use. Must support set_sample_schema().
+    :param logger: Logger to use.
+    :param name: Name of the indexer for logging.
+    :param sample_generator_kwargs: Other kwargs passed to TokenizedSampleWithAuxGenerator
+        (e.g., sample_size, token_dtype).
+
+    :raises ValueError: If both input_base_path and input_file_paths are None.
+    :raises ValueError: If output destination cannot be determined.
+
+    :return: FileTextLinesIndexer instance that can be called to run the pipeline.
+    """
+    # Import here to avoid circular dependency
+    from data_forager.sample_generators.tokenization_with_aux import (
+        TokenizedSampleWithAuxGenerator,
+    )
+
+    if logger is None:
+        logger = module_logger
+
+    # Validate input source
+    if input_base_path is None and input_file_paths is None:
+        raise ValueError(
+            "Either input_base_path or input_file_paths must be provided"
+        )
+
+    # Determine output base path
+    effective_output_base_path = output_base_path or input_base_path
+
+    # Validate output destination
+    if index_store is None and effective_output_base_path is None:
+        raise ValueError(
+            "Either index_store, output_base_path, or input_base_path must be provided "
+            "to determine where to store the index"
+        )
+
+    logger.info(f"Output base path: {effective_output_base_path}")
+
+    if index_store is None:
+        index_store = FSBasedIndexStore(
+            base_path=effective_output_base_path,
+        )
+
+    if input_file_paths is None:
+        logger.info(f"Scanning for JSONL files in: {input_base_path}")
+        input_file_paths = find_files_recursive(
+            input_base_path,
+            extension_patterns=['*.jsonl', '*.JSONL']
+        )
+        input_file_paths = natural_sort(input_file_paths)
+        logger.info(f"Found {len(input_file_paths)} JSONL file(s)")
+
+    # Set default base_output_path for tokenized samples if not provided
+    if 'base_output_path' not in sample_generator_kwargs:
+        default_base_output_path = os.path.join(
+            effective_output_base_path, "tokenized-samples"
+        )
+        logger.info(f"Tokenized samples output path: {default_base_output_path}")
+        sample_generator_kwargs['base_output_path'] = default_base_output_path
+
+    sample_generator = TokenizedSampleWithAuxGenerator(
+        process_parts_func=process_parts_func,
+        tokenizer_func=tokenizer_func,
+        eos_idx=eos_idx,
+        aux_generators=aux_generators,
+        **sample_generator_kwargs
+    )
+
+    return FileTextLinesIndexer(
+        input_file_paths=input_file_paths,
+        index_store=index_store,
+        sample_generator=sample_generator,
+        description="Tokenizing with aux data and indexing",
+        name=name,
+    )

data_forager-0.2.0/data_forager/sample_generators/__init__.py

@@ -0,0 +1,30 @@
+"""
+Sample generators for transforming input data into samples.
+
+This package contains:
+- SampleGeneratorInterface: Protocol for sample generators
+- SampleData: Data class for sample information
+- SampleSchema, ArraySpec: Schema classes for structured samples
+- TokenizedSampleGenerator: Tokenizes text into fixed-length samples
+- TokenizedSampleWithAuxGenerator: Tokenizes with auxiliary data (loss masks, etc.)
+"""
+
+from data_forager.sample_generators.common import (
+    SampleData,
+    SampleGeneratorInterface,
+    NOOPSampleGenerator,
+    noop_sample_processing,
+)
+from data_forager.sample_generators.schema import (
+    ArraySpec,
+    SampleSchema,
+)
+
+__all__ = [
+    "ArraySpec",
+    "NOOPSampleGenerator",
+    "SampleData",
+    "SampleGeneratorInterface",
+    "SampleSchema",
+    "noop_sample_processing",
+]