data-forager 0.1.1__py3-none-any.whl → 0.1.3__py3-none-any.whl

data_forager/index_stores/common.py

@@ -6,16 +6,25 @@ from data_forager.sample_index import SampleIndex
 class IndexStoreInterface(Protocol):
 
     def init_store(self):
+        """Initialize the store for writing. Must be called before add_sample()."""
         ...
 
     def add_sample(self, file_location: str, byte_offset: int, num_bytes: int):
+        """Add a sample location to the index."""
+        ...
+
+    def close(self):
+        """Close the store, flushing any buffered data. Must be called after all samples are added."""
         ...
 
     def load(self) -> SampleIndex:
+        """Load the index from the store."""
         ...
 
     def exists(self) -> bool:
+        """Check if the index already exists."""
         ...
 
     def clear(self) -> None:
+        """Remove the index."""
         ...

data_forager/index_stores/fs_based.py

@@ -1,6 +1,6 @@
 import os.path
 import shutil
-from typing import Optional
+from typing import Optional, TextIO, BinaryIO
 
 import numpy as np
 
@@ -10,7 +10,7 @@ from data_forager.index_stores.common import IndexStoreInterface
 from data_forager.sample_index import SampleIndex
 
 
-class IndexStore(IndexStoreInterface, Base):
+class IndexStore(Base, IndexStoreInterface):
 
     def __init__(self, base_path: str, index_data_folder: str = "index", name: Optional[str] = None):
         """
@@ -25,12 +25,24 @@ class IndexStore(IndexStoreInterface, Base):
         self._index_data_path = os.path.join(base_path, index_data_folder)
         self._file_locations = []
 
+        # File handles for buffered writing
+        self._file_location_handle: Optional[TextIO] = None
+        self._sample_locations_handle: Optional[BinaryIO] = 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}")
 
         os.makedirs(self._index_data_path, exist_ok=True)
 
+        # Open file handles for writing
+        self._file_location_handle = open(
+            os.path.join(self._index_data_path, "file_location.txt"), "a"
+        )
+        self._sample_locations_handle = open(
+            os.path.join(self._index_data_path, "sample_locations.bin"), "ab"
+        )
+
     def add_sample(self, file_location: str, byte_offset: int, num_bytes: int):
         """
         :param file_location:
@@ -40,15 +52,25 @@ class IndexStore(IndexStoreInterface, Base):
         """
         if file_location not in self._file_locations:
             self._file_locations.append(file_location)
-            with open(os.path.join(self._index_data_path, "file_location.txt"), "a") as f:
-                f.writelines([file_location+'\n'])
+            self._file_location_handle.write(file_location + '\n')
 
         file_index = self._file_locations.index(file_location)
 
-        with open(os.path.join(self._index_data_path, "sample_locations.bin"), "ab") as f:
-            sample_location_bytes = np.array([file_index, byte_offset, num_bytes], dtype=np.uint64).tobytes()
+        sample_location_bytes = np.array([file_index, byte_offset, num_bytes], dtype=np.uint64).tobytes()
+        self._sample_locations_handle.write(sample_location_bytes)
+
+    def close(self):
+        """Close file handles and flush buffered data."""
+        if self._file_location_handle is not None:
+            self._file_location_handle.close()
+            self._file_location_handle = None
+
+        if self._sample_locations_handle is not None:
+            self._sample_locations_handle.close()
+            self._sample_locations_handle = None
 
-            f.write(sample_location_bytes)
+    def __del__(self):
+        self.close()
 
     def load(self) -> SampleIndex:
         with open(os.path.join(self._index_data_path, "file_location.txt"), "r") as f:

data_forager/indexers/text_lines.py

@@ -155,3 +155,5 @@ class FileTextLinesIndexer(Base):
 
             sys.stdout.write('\n\n')
             sys.stdout.flush()
+
+        self._index_store.close()

data_forager/indexers/tokenization_indexer.py

@@ -1,5 +1,6 @@
 from typing import Callable, List, Optional
 
+import logging
 import os
 from pathlib import Path
 
@@ -8,7 +9,11 @@ import json
 import numpy as np
 
 from basics.base import Base
+from basics.logging import get_logger
 
+module_logger = get_logger(os.path.basename(__file__))
+
+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 SampleData, FileTextLinesIndexer, SampleGeneratorInterface
 from data_forager.utils import find_files_recursive, natural_sort
@@ -24,61 +29,125 @@ def get_text_from_jsonl(jsonl_bytes: bytes, text_key: str = "text", text_encodin
 
 
 def create_tokenize_and_index_jsonl_text_func(
-    input_base_path: str,
     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 function to:
-     * Tokenize text from input JSONL objects, loaded from files at input_base_path (recursively),
-     * Store the token data in bin files under folder "tokenized-samples" in input_base_path
-     * Store index data under folder "index" in input_base_path
+    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:
-    # Create pipeline to tokenize text from input JSONL objects and index the token samples
+    ```python
     import tiktoken
 
     enc = tiktoken.get_encoding("gpt2")
     def tokenize_text(text: str) -> List[int]:
-        return tiktoken.enc.encode_ordinary(text)
+        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,
+    )
 
-    tokenize_and_index_jsonl_text_func = create_jsonl_text_tokenization_and_indexing_pipeline(
-        input_base_path='.',
+    # Option 2: Explicit input files and output path
+    indexer = create_tokenize_and_index_jsonl_text_func(
         tokenizer_func=tokenize_text,
-        sample_size=1024
+        eos_idx=enc.eot_token,
+        input_file_paths=['./data/train.jsonl'],
+        output_base_path='./output',
+        sample_size=1024,
     )
 
-    # Start tokenization and indexing
-    tokenize_and_index_jsonl_text_func()
+    # Run tokenization and indexing
+    indexer()
+    ```
 
-    :param input_base_path: Path to directory containing JSONL files (searched recursively).
     :param tokenizer_func: Function used to tokenize text.
-    :param eos_idx: EOS token index, known by the used Tokenizer
+    :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 sample_generator_kwargs: Other kwargs passed to TokenizedSampleGenerator.
-    :param name: Optional: name of the indexer to create, used for logging purposes
+    :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 used to tokenize and index text from jsonl objects, from
-             JSONL files at input_base_path (recursively)
+    :return: FileTextLinesIndexer instance that can be called to run tokenization
+        and indexing.
     """
-    if process_text_line_func is None:
-        process_text_line_func=get_text_from_jsonl
+    if logger is None:
+        logger = module_logger
 
-    index_store = FSBasedIndexStore(
-        base_path=input_base_path,
-    )
-    input_file_paths = find_files_recursive(
-        input_base_path,
-        extension_patterns=['*.jsonl', '*.JSONL']
-    )
+    # 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
 
-    # Assuming numbered files
-    input_file_paths = natural_sort(input_file_paths)
+    # 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,
@@ -96,7 +165,7 @@ def create_tokenize_and_index_jsonl_text_func(
     )
 
 
-class TokenizedSampleGenerator(SampleGeneratorInterface, Base):
+class TokenizedSampleGenerator(Base, SampleGeneratorInterface):
 
     def __init__(
         self,
@@ -215,6 +284,10 @@ class TokenizedSampleGenerator(SampleGeneratorInterface, Base):
         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
@@ -224,8 +297,9 @@ class TokenizedSampleGenerator(SampleGeneratorInterface, Base):
             num_rest_tokens = num_tokens % self._sample_size
 
             if num_rest_tokens > 0:
-                self._rest_tokens = tokenized_text[-num_rest_tokens:] + [self._eos_idx]
-                tokenized_text = tokenized_text[:num_samples*self._sample_size]
+                # 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)

{data_forager-0.1.1.dist-info → data_forager-0.1.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: data-forager
-Version: 0.1.1
+Version: 0.1.3
 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.1.1.dist-info → data_forager-0.1.3.dist-info}/RECORD

@@ -7,14 +7,14 @@ data_forager/datasets/common.py,sha256=OB0DzuxuLfd9BP68K9vTNzzBfK-L6ZzIFqe-FZTwQ
 data_forager/datasets/jsonl.py,sha256=enOjWRT-AJTF3tWtNlonCqdDpZfVsK8If6yEtlA8tns,630
 data_forager/datasets/tokens.py,sha256=OP5MNb9uBDSX_Of6lNVLs5CAj46RwkP4gGDk-94lD40,597
 data_forager/index_stores/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-data_forager/index_stores/common.py,sha256=4ak_LGOWp1jhRVkjyxuYPQ7vf7qNiSIfJB6UwKl-X0I,383
-data_forager/index_stores/fs_based.py,sha256=apwSQegacadBr9GB1amcelz8TseeAahYNA1cNv1gCjs,2798
+data_forager/index_stores/common.py,sha256=HYPZfCtmbLDxjPgCoUseahicTEI1el7V1s5tSyZxkfs,780
+data_forager/index_stores/fs_based.py,sha256=1MB_iZPZS187rlLRUNPIJcQK-lU9EqEAg3sARnT2vSk,3564
 data_forager/indexers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 data_forager/indexers/jsonl_indexer.py,sha256=xvogFjEcKPNr-GBkKEr8WnWPVvJbkWQzUBrGwVELlA4,1395
-data_forager/indexers/text_lines.py,sha256=Z5XmlhfR1DiLAF_2GUO88H5SJFOhUGwQ0vkcGfltRXg,5058
-data_forager/indexers/tokenization_indexer.py,sha256=gzAuDwIJenNklfHkegDoZz-JzIOXmvuwj4r10wtbJpo,10908
-data_forager-0.1.1.dist-info/licenses/LICENSE,sha256=If0vYAiJJUtbASoyZPVhvTu3e3m4WB1cQmUpvo9HRTc,1071
-data_forager-0.1.1.dist-info/METADATA,sha256=P_ZlNLS3x9O2OkPyNLgsB58k_jooKMlnpDYkh7l9x9w,9090
-data_forager-0.1.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-data_forager-0.1.1.dist-info/top_level.txt,sha256=mLcF2mYnfdaeJ_vIa7hT-MtTpUvn7kgyaWNuxXZ1Ds8,13
-data_forager-0.1.1.dist-info/RECORD,,
+data_forager/indexers/text_lines.py,sha256=XMm5oc0btP7I16z87g1fmq9AqJyVhDOvR2cDu_zFZio,5093
+data_forager/indexers/tokenization_indexer.py,sha256=t-7Q3PLAJ0DYZT6LWdHeahk9Hz9OQsvWfoPvhHIneMk,13927
+data_forager-0.1.3.dist-info/licenses/LICENSE,sha256=If0vYAiJJUtbASoyZPVhvTu3e3m4WB1cQmUpvo9HRTc,1071
+data_forager-0.1.3.dist-info/METADATA,sha256=S8OWa1JSkQB77oVn3QMVEPN6obnvePn12jlojzZvELo,9090
+data_forager-0.1.3.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+data_forager-0.1.3.dist-info/top_level.txt,sha256=mLcF2mYnfdaeJ_vIa7hT-MtTpUvn7kgyaWNuxXZ1Ds8,13
+data_forager-0.1.3.dist-info/RECORD,,