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

data_forager/sample_generators/common.py

@@ -0,0 +1,117 @@
+"""
+Common interfaces and utilities for sample generators.
+
+This module provides the base protocol and data structures used by all sample generators.
+"""
+
+from typing import List, Protocol, TYPE_CHECKING
+
+from dataclasses import dataclass
+
+if TYPE_CHECKING:
+    from data_forager.sample_generators.schema import SampleSchema
+
+
+@dataclass
+class SampleData:
+    """
+    Data class containing sample bytes and file location.
+
+    :param sample_bytes: The sample data as bytes.
+    :param file_path: Path to the file where the sample is stored.
+    """
+
+    sample_bytes: bytes
+    file_path: str
+
+
+class SampleGeneratorInterface(Protocol):
+    """
+    Protocol defining the interface for sample generators.
+
+    Sample generators transform input data (e.g., text lines) into samples
+    that can be indexed and stored for random access during training.
+    """
+
+    def prepare(self, text_file_path: str):
+        """
+        Prepare sample generation from a new input text file.
+
+        :param text_file_path: Path to the input text file.
+        """
+        ...
+
+    def create_samples(self, text_line: bytes) -> List[SampleData]:
+        """
+        Create one or more samples from the given text_line and store them.
+
+        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 SampleData objects. For each created sample:
+            - 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 the current input file.
+
+        Called after all text lines from the input file have been processed.
+
+        :param is_last_file: Indicates if the input text file was the last file
+            to be processed.
+        """
+        ...
+
+    def get_sample_schema(self) -> "SampleSchema | None":
+        """
+        Return schema describing sample structure, or None for unstructured samples.
+
+        The schema describes how sample bytes should be interpreted (e.g., array
+        names, dtypes, and offsets for multi-array samples).
+
+        :return: SampleSchema if samples have structured format, None otherwise.
+        """
+        ...
+
+
+class NOOPSampleGenerator:
+    """
+    A no-operation sample generator that returns input lines unchanged.
+
+    Used as the default generator when no transformation is needed.
+    """
+
+    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 get_sample_schema(self) -> "SampleSchema | None":
+        return None
+
+
+def noop_sample_processing(text_line: bytes, text_file_path: str) -> List[SampleData]:
+    """
+    Simple function that wraps a text line as a sample without transformation.
+
+    :param text_line: The input text line as bytes.
+    :param text_file_path: The file path to associate with the sample.
+
+    :return: List containing a single SampleData with the unchanged text line.
+    """
+    return [SampleData(text_line, text_file_path)]

data_forager/sample_generators/schema.py

@@ -0,0 +1,54 @@
+"""
+Schema definitions for describing sample structure.
+
+This module contains dataclasses that describe how sample bytes should be
+interpreted, particularly for samples containing multiple arrays (e.g., tokens
+and loss masks).
+"""
+
+from typing import List, Optional
+
+from dataclasses import dataclass, field
+
+
+@dataclass
+class ArraySpec:
+    """
+    Specification for a single array within a sample.
+
+    :param name: Name of the array (e.g., "tokens", "loss_mask").
+    :param dtype: NumPy dtype string (e.g., "uint32", "uint8").
+    :param offset: Byte offset within the sample where this array starts.
+    """
+
+    name: str
+    dtype: str
+    offset: int
+
+
+@dataclass
+class SampleSchema:
+    """
+    Schema describing the structure of samples.
+
+    Used when samples contain multiple arrays (e.g., tokens + auxiliary data).
+    Each array has a name, dtype, and byte offset within the sample.
+
+    :param sample_size: Number of elements per array (e.g., context length).
+    :param arrays: List of ArraySpec describing each array in the sample.
+    :param total_bytes_per_sample: Total size of each sample in bytes.
+        If None, computed automatically from arrays.
+    """
+
+    sample_size: int
+    arrays: List[ArraySpec] = field(default_factory=list)
+    total_bytes_per_sample: Optional[int] = None
+
+    def __post_init__(self):
+        """Compute total_bytes_per_sample if not provided."""
+        if self.total_bytes_per_sample is None and self.arrays:
+            import numpy as np
+            total = 0
+            for arr in self.arrays:
+                total += self.sample_size * np.dtype(arr.dtype).itemsize
+            self.total_bytes_per_sample = total

data_forager/sample_generators/tokenization.py

@@ -0,0 +1,210 @@
+"""
+Tokenized sample generator for converting text into fixed-length token samples.
+
+This module provides TokenizedSampleGenerator which tokenizes text and produces
+samples of a fixed context length, suitable for language model training.
+"""
+
+from typing import Callable, List, Optional
+
+import os
+from pathlib import Path
+
+import numpy as np
+
+from basics.base import Base
+
+from data_forager.sample_generators.common import SampleData, SampleGeneratorInterface
+
+
+TokenizerFunc = Callable[[str], List[int]]
+ProcessTextLineFunc = Callable[[bytes], str]
+
+
+class TokenizedSampleGenerator(Base, SampleGeneratorInterface):
+    """
+    Tokenizes text into fixed-length samples for language model training.
+
+    This generator:
+    1. Processes input text lines using a configurable function
+    2. Tokenizes the text using a provided tokenizer
+    3. Splits tokens into fixed-length samples
+    4. Handles document boundaries with EOS tokens
+    5. Carries over remainder tokens across documents
+    """
+
+    def __init__(
+        self,
+        process_text_line_func: ProcessTextLineFunc,
+        tokenizer_func: TokenizerFunc,
+        eos_idx: int,
+        token_dtype: np.dtype = np.uint16,
+        sample_size: Optional[int] = None,
+        base_output_path: str = None,
+        file_name_postfix: str = "tokenized-samples",
+        name: Optional[str] = None
+    ):
+        """
+        Initialize the tokenized sample generator.
+
+        Tokenizes and indexes text into fixed length (`sample_size` not None) samples or
+        samples of variable size, depending on the text (`sample_size` is None).
+
+        This callable performs the following steps:
+
+        ## prepare ##
+        * In the preparation step, create file to store tokenized samples, based on the
+          input `text_file_path` and the given `base_output_path` and `file_name_postfix`
+        * If the `base_output_path` is not given the `text_file_path` will be used +
+          "/tokenized-samples"
+
+        ## create_samples ##
+        * To create tokenized text samples, processes incoming text line using
+          `process_text_line_func`, e.g. convert JSONL into dict and retrieve the sample
+          text from it.
+        * The resulting text is tokenized using `tokenizer_func`.
+        * If a `sample_size` is given:
+          The tokenized text is split into samples of length `sample_size` and stored in
+          the file opened in the prepare step. Here `token_dtype` is used.
+          - Trailing tokens will be combined with samples of a next text line
+          - Tokens from different text samples will be separated by `eos_idx`
+        * If a `sample_size` is not given:
+          The tokenized text is immediately stored as is, in the file opened in the
+          prepare step. Here `token_dtype` is used.
+
+        ## finish ##
+        * After all text lines are processed, the file holding the tokenized text samples
+          is closed. When `sample_size` is not None: Any final trailing tokens will be
+          discarded, but only when the last input text file was processed.
+
+        :param process_text_line_func: Function to extract text from input bytes.
+        :param tokenizer_func: Function to tokenize text into token IDs.
+        :param eos_idx: End-of-sequence token ID for document boundaries.
+        :param token_dtype: NumPy dtype for storing tokens (default: uint16).
+        :param sample_size: Fixed sample length, or None for variable-length samples.
+        :param base_output_path: Base path for output files.
+        :param file_name_postfix: Postfix for output file names.
+        :param name: Name for logging purposes.
+        """
+        super().__init__(pybase_logger_name=name)
+
+        if sample_size is None:
+            self._log.info("Tokenized text will NOT be broken into samples of fixed length.")
+
+        self._process_text_line_func = process_text_line_func
+        self._tokenizer_func = tokenizer_func
+        self._eos_idx = eos_idx
+        self._token_dtype = token_dtype
+        self._sample_size = sample_size
+        self._base_output_path = base_output_path
+        self._file_name_postfix = file_name_postfix
+
+        self._current_samples_path = None
+        self._current_samples_file = None
+
+        self._rest_tokens = None
+
+    def prepare(self, text_file_path: str):
+        """
+        Prepare for processing a new input file.
+
+        Creates the output file for storing tokenized samples based on the input
+        `text_file_path` and the configured `base_output_path` and `file_name_postfix`.
+
+        :param text_file_path: Path to the input text file.
+        """
+        input_file_path = os.path.dirname(text_file_path)
+        input_file_name = Path(text_file_path).stem
+        output_file_name = f"{input_file_name}-{self._file_name_postfix}.bin"
+
+        output_path = self._base_output_path
+        if self._base_output_path is None:
+            output_path = os.path.join(input_file_path, "tokenized-samples")
+
+        os.makedirs(output_path, exist_ok=True)
+
+        output_file_path = os.path.join(output_path, output_file_name)
+        if os.path.exists(output_file_path):
+            raise FileExistsError(f"Tokenized samples file already exists: \n{output_file_path}")
+
+        self._current_samples_path = output_file_path
+        self._current_samples_file = open(output_file_path, "wb")
+
+        self._log.debug(f"Tokenized samples file opened: \n{output_file_path}")
+
+    def create_samples(self, text_line: bytes) -> List[SampleData]:
+        """
+        Create tokenized samples from a text line.
+
+        Processes the input text line, tokenizes it, and splits into fixed-length
+        samples (if sample_size is set). Handles document boundaries with EOS tokens
+        and carries over remainder tokens.
+
+        :param text_line: JSONL text line as bytes.
+
+        :return: List of SampleData objects containing the tokenized samples.
+        """
+        input_text = self._process_text_line_func(text_line)
+        tokenized_text = self._tokenizer_func(input_text)
+
+        if self._sample_size is not None:
+            # Always append EOS after each document to mark document boundary
+            tokenized_text = tokenized_text + [self._eos_idx]
+
+            # Prepend any leftover tokens from previous document
+            if self._rest_tokens is not None:
+                tokenized_text = self._rest_tokens + tokenized_text
+                self._rest_tokens = None
+
+            num_tokens = len(tokenized_text)
+            num_samples = num_tokens // self._sample_size
+            num_rest_tokens = num_tokens % self._sample_size
+
+            if num_rest_tokens > 0:
+                # Store remainder tokens (includes EOS from this document)
+                self._rest_tokens = tokenized_text[-num_rest_tokens:]
+                tokenized_text = tokenized_text[:num_samples * self._sample_size]
+
+            tokenized_samples = np.array(tokenized_text, dtype=self._token_dtype)
+            tokenized_samples = tokenized_samples.reshape(-1, self._sample_size)
+        else:
+            tokenized_samples = np.array([tokenized_text], dtype=self._token_dtype)
+
+        # Store tokenized_samples
+        sample_data = []
+        for sample_idx in range(tokenized_samples.shape[0]):
+            sample_bytes = tokenized_samples[sample_idx, :].tobytes()
+            sample_data.append(SampleData(
+                sample_bytes, self._current_samples_path,
+            ))
+
+            self._current_samples_file.write(sample_bytes)
+
+        return sample_data
+
+    def finish(self, is_last_file: bool):
+        """
+        Finish processing the current input file.
+
+        Closes the current samples file. When `sample_size` is not None, any final
+        trailing tokens will be discarded only when the last input file was processed.
+
+        :param is_last_file: Whether this is the last input file.
+        """
+        self._close_current_samples_file()
+
+        if is_last_file and self._rest_tokens is not None:
+            self._log.debug(f"Cut off {len(self._rest_tokens)} unused tokens")
+
+    def get_sample_schema(self):
+        """Return None as this generator produces simple token arrays."""
+        return None
+
+    def _close_current_samples_file(self):
+        if self._current_samples_file:
+            self._log.debug(f"Closing tokenized samples file: \n{self._current_samples_path}")
+            self._current_samples_file.close()
+            self._current_samples_file = None
+
+    def __del__(self):
+        self._close_current_samples_file()

data_forager/sample_generators/tokenization_with_aux.py

@@ -0,0 +1,250 @@
+"""
+Tokenized sample generator with auxiliary data support.
+
+This module provides TokenizedSampleWithAuxGenerator which tokenizes structured
+samples (with typed parts) and generates auxiliary data (e.g., loss masks)
+alongside the tokens.
+"""
+
+from typing import Callable, Dict, List, Optional
+
+import os
+from pathlib import Path
+
+import numpy as np
+
+from basics.base import Base
+
+from data_forager.sample_generators.common import SampleData, SampleGeneratorInterface
+from data_forager.sample_generators.schema import SampleSchema, ArraySpec
+from data_forager.sample_generators.aux.common import AuxDataGenerator, Part
+
+
+ProcessPartsFunc = Callable[[bytes], List[Part]]
+TokenizerFunc = Callable[[str], List[int]]
+
+
+class TokenizedSampleWithAuxGenerator(Base, SampleGeneratorInterface):
+    """
+    Tokenizes structured samples and generates auxiliary data.
+
+    This generator handles samples with typed parts (e.g., prompt, response)
+    and produces both tokens and auxiliary data (e.g., loss masks) in a single
+    concatenated output.
+
+    Input format: JSONL with {"parts": [{"type": "...", "text": "..."}, ...]}
+    Output format: Concatenated bytes [tokens][aux1][aux2]...
+    """
+
+    def __init__(
+        self,
+        process_parts_func: ProcessPartsFunc,
+        tokenizer_func: TokenizerFunc,
+        eos_idx: int,
+        aux_generators: Dict[str, AuxDataGenerator],
+        token_dtype: np.dtype = np.uint32,
+        sample_size: int = 4096,
+        base_output_path: Optional[str] = None,
+        file_name_postfix: str = "tokenized-samples",
+        name: Optional[str] = None,
+    ):
+        """
+        Initialize the tokenized sample generator with auxiliary data support.
+
+        :param process_parts_func: Function to extract typed parts from input bytes.
+            Takes JSONL bytes and returns List[Part].
+        :param tokenizer_func: Function to tokenize text into token IDs.
+        :param eos_idx: End-of-sequence token ID for document boundaries.
+        :param aux_generators: Dict mapping names to AuxDataGenerator instances.
+            Example: {"loss_mask": LossMaskGenerator()}
+        :param token_dtype: NumPy dtype for storing tokens (default: uint32).
+        :param sample_size: Fixed sample length in tokens.
+        :param base_output_path: Base path for output files.
+        :param file_name_postfix: Postfix for output file names.
+        :param name: Name for logging purposes.
+        """
+        super().__init__(pybase_logger_name=name)
+
+        if not aux_generators:
+            raise ValueError("aux_generators must not be empty")
+
+        self._process_parts_func = process_parts_func
+        self._tokenizer_func = tokenizer_func
+        self._eos_idx = eos_idx
+        self._aux_generators = aux_generators
+        self._token_dtype = np.dtype(token_dtype)
+        self._sample_size = sample_size
+        self._base_output_path = base_output_path
+        self._file_name_postfix = file_name_postfix
+
+        # Build schema
+        self._sample_schema = self._build_schema()
+
+        # Current output state
+        self._current_samples_path: Optional[str] = None
+        self._current_samples_file = None
+
+        # Remainder from previous document
+        self._rest_tokens: Optional[List[int]] = None
+        self._rest_aux: Optional[Dict[str, List[int]]] = None
+
+    def _build_schema(self) -> SampleSchema:
+        """Build the sample schema from token dtype and aux generators."""
+        arrays = []
+        offset = 0
+
+        # Tokens array first
+        token_bytes = self._sample_size * self._token_dtype.itemsize
+        arrays.append(ArraySpec(name="tokens", dtype=str(self._token_dtype), offset=offset))
+        offset += token_bytes
+
+        # Auxiliary arrays in sorted order (for deterministic output)
+        for name in sorted(self._aux_generators.keys()):
+            gen = self._aux_generators[name]
+            aux_bytes = self._sample_size * gen.dtype.itemsize
+            arrays.append(ArraySpec(name=name, dtype=str(gen.dtype), offset=offset))
+            offset += aux_bytes
+
+        return SampleSchema(
+            sample_size=self._sample_size,
+            arrays=arrays,
+            total_bytes_per_sample=offset,
+        )
+
+    def get_sample_schema(self) -> SampleSchema:
+        """Return the sample schema describing the output structure."""
+        return self._sample_schema
+
+    def prepare(self, text_file_path: str):
+        """
+        Prepare for processing a new input file.
+
+        Creates the output file for storing tokenized samples with aux data.
+
+        :param text_file_path: Path to the input text file.
+        """
+        input_file_path = os.path.dirname(text_file_path)
+        input_file_name = Path(text_file_path).stem
+        output_file_name = f"{input_file_name}-{self._file_name_postfix}.bin"
+
+        output_path = self._base_output_path
+        if self._base_output_path is None:
+            output_path = os.path.join(input_file_path, "tokenized-samples")
+
+        os.makedirs(output_path, exist_ok=True)
+
+        output_file_path = os.path.join(output_path, output_file_name)
+        if os.path.exists(output_file_path):
+            raise FileExistsError(f"Tokenized samples file already exists: \n{output_file_path}")
+
+        self._current_samples_path = output_file_path
+        self._current_samples_file = open(output_file_path, "wb")
+
+        self._log.debug(f"Tokenized samples file opened: \n{output_file_path}")
+
+    def create_samples(self, text_line: bytes) -> List[SampleData]:
+        """
+        Create tokenized samples with auxiliary data from a structured input.
+
+        Processes the input parts, tokenizes each part, generates auxiliary data,
+        and splits into fixed-length samples.
+
+        :param text_line: JSONL line as bytes containing structured parts.
+
+        :return: List of SampleData objects containing the tokenized samples.
+        """
+        parts = self._process_parts_func(text_line)
+
+        # Tokenize all parts and generate aux data
+        doc_tokens: List[int] = []
+        doc_aux: Dict[str, List[int]] = {name: [] for name in self._aux_generators}
+
+        for part in parts:
+            part_tokens = self._tokenizer_func(part.text)
+            doc_tokens.extend(part_tokens)
+
+            for name, gen in self._aux_generators.items():
+                aux_values = gen.generate(
+                    part_type=part.type,
+                    num_tokens=len(part_tokens),
+                    part_tokens=part_tokens,
+                )
+                doc_aux[name].extend(aux_values)
+
+        # Add EOS token and aux values
+        doc_tokens.append(self._eos_idx)
+        for name, gen in self._aux_generators.items():
+            doc_aux[name].append(gen.generate_for_eos())
+
+        # Prepend any leftover from previous document
+        # Invariant: rest_tokens and all rest_aux[*] have the same length
+        if self._rest_tokens is not None:
+            assert all(
+                len(self._rest_aux[name]) == len(self._rest_tokens)
+                for name in self._aux_generators
+            ), "Invariant violated: rest_tokens and rest_aux must have same length"
+
+            doc_tokens = self._rest_tokens + doc_tokens
+            for name in self._aux_generators:
+                doc_aux[name] = self._rest_aux[name] + doc_aux[name]
+            self._rest_tokens = None
+            self._rest_aux = None
+
+        # Split into samples
+        num_tokens = len(doc_tokens)
+        num_samples = num_tokens // self._sample_size
+        num_rest_tokens = num_tokens % self._sample_size
+
+        if num_rest_tokens > 0:
+            # Store remainder
+            self._rest_tokens = doc_tokens[-num_rest_tokens:]
+            self._rest_aux = {name: doc_aux[name][-num_rest_tokens:] for name in self._aux_generators}
+            doc_tokens = doc_tokens[:num_samples * self._sample_size]
+            for name in self._aux_generators:
+                doc_aux[name] = doc_aux[name][:num_samples * self._sample_size]
+
+        # Write samples
+        sample_data = []
+        for i in range(num_samples):
+            start = i * self._sample_size
+            end = start + self._sample_size
+
+            # Build concatenated sample bytes
+            sample_tokens = doc_tokens[start:end]
+            sample_bytes = np.array(sample_tokens, dtype=self._token_dtype).tobytes()
+
+            # Add aux data in sorted order
+            for name in sorted(self._aux_generators.keys()):
+                gen = self._aux_generators[name]
+                aux_values = doc_aux[name][start:end]
+                sample_bytes += np.array(aux_values, dtype=gen.dtype).tobytes()
+
+            sample_data.append(SampleData(sample_bytes, self._current_samples_path))
+            self._current_samples_file.write(sample_bytes)
+
+        return sample_data
+
+    def finish(self, is_last_file: bool):
+        """
+        Finish processing the current input file.
+
+        Closes the current samples file. When is_last_file is True, any
+        remaining tokens are discarded.
+
+        :param is_last_file: Whether this is the last input file.
+        """
+        self._close_current_samples_file()
+
+        if is_last_file and self._rest_tokens is not None:
+            self._log.debug(f"Cut off {len(self._rest_tokens)} unused tokens")
+            self._rest_tokens = None
+            self._rest_aux = None
+
+    def _close_current_samples_file(self):
+        if self._current_samples_file:
+            self._log.debug(f"Closing tokenized samples file: \n{self._current_samples_path}")
+            self._current_samples_file.close()
+            self._current_samples_file = None
+
+    def __del__(self):
+        self._close_current_samples_file()

data_forager/sample_index.py

@@ -1,12 +1,29 @@
-from typing import List, Optional, Dict, Protocol
+"""
+Sample index data structures for random access to samples.
 
-from dataclasses import dataclass
+This module provides the core data structures for indexing samples:
+- SampleLocation: Location of a single sample (file path, byte offset, size)
+- SampleIndex: Index containing all sample locations and optional schema
+"""
+
+from typing import List, Optional
+
+from dataclasses import dataclass, field
 
 import numpy as np
 
+from data_forager.sample_generators.schema import SampleSchema
+
 
 @dataclass
 class SampleLocation:
+    """
+    Location information for a single sample.
+
+    :param file_path: Path to the file containing the sample.
+    :param byte_offset: Byte offset within the file where the sample starts.
+    :param num_bytes: Size of the sample in bytes.
+    """
 
     file_path: str
     byte_offset: int
@@ -20,6 +37,18 @@ class SampleLocation:
 
 @dataclass
 class SampleIndex:
+    """
+    Index for random access to samples stored on disk.
+
+    Contains file locations and byte offsets for all samples, enabling O(1)
+    random access via seek operations.
+
+    :param file_locations: List of file paths containing samples.
+    :param sample_locations: 2D numpy array of shape (num_samples, 3) where
+        each row contains (file_index, byte_offset, num_bytes).
+    :param sample_schema: Optional schema describing sample structure for
+        samples with auxiliary data (e.g., tokens + loss_mask).
+    """
 
     file_locations: List[str]
 
@@ -30,6 +59,9 @@ class SampleIndex:
     #     num_bytes  : uint64
     sample_locations: np.ndarray
 
+    # Optional schema for structured samples (e.g., tokens + auxiliary data)
+    sample_schema: Optional[SampleSchema] = field(default=None)
+
     def __len__(self) -> int:
         return self.sample_locations.shape[0]