feat: deterministic streaming dataset tokenization

openverifiablellm/tokenizer/__init__.py

@@ -4,3 +4,5 @@
     "train_tokenizer",
     "hash_tokenizer_config",
 ]
+
+from .tokenize_dataset import tokenize_dataset

openverifiablellm/tokenizer/tokenize_dataset.py

@@ -0,0 +1,62 @@
+from pathlib import Path
+import numpy as np
+
+
+def tokenize_dataset(input_file, tokenizer, output_file):
+    """
+    Tokenize a dataset using a trained tokenizer and save tokens to a binary file.
+
+    This implementation is streaming and memory-efficient, meaning it can handle
+    very large datasets without loading everything into memory.
+
+    Parameters
+    ----------
+    input_file : str or Path
+        Path to the cleaned dataset text file.
+
+    tokenizer : object
+        A tokenizer instance with an `encode()` method.
+
+    output_file : str or Path
+        Path where tokenized binary output will be written.
+
+    Returns
+    -------
+    int
+        Total number of tokens written.
+    """
+
+    input_path = Path(input_file)
+    output_path = Path(output_file)
+
+    if not input_path.exists():
+        raise FileNotFoundError(f"Dataset file not found: {input_path}")
+
+    total_tokens = 0
+
+    # open dataset for streaming
+    with input_path.open("r", encoding="utf-8") as fin, \
+         output_path.open("wb") as fout:
+
+        for line in fin:
+
+            text = line.strip()
+
+            if not text:
+                continue
+
+            encoded = tokenizer.encode(text)
+
+            if isinstance(encoded, list):
+                tokens = encoded
+            else:
+                # support tokenizers returning objects
+                tokens = encoded.ids
+
+            tokens_array = np.array(tokens, dtype=np.uint32)
+
+            tokens_array.tofile(fout)
+
+            total_tokens += len(tokens)
+
+    return total_tokens

pyproject.toml

@@ -14,7 +14,8 @@ requires-python = ">=3.9"
 dependencies = [
     "defusedxml",
     "sentencepiece",
-    "tokenizers==0.15.2"
+    "tokenizers==0.15.2",
+    "numpy"
 ]
 
 [tool.setuptools.packages.find]

tests/test_tokenizer.py

@@ -1,7 +1,9 @@
 import json
 
 import pytest
-
+from pathlib import Path
+import tempfile
+import numpy as np
 from openverifiablellm.tokenizer import (
     hash_tokenizer_config,
     train_tokenizer,
@@ -166,3 +168,47 @@ def test_hash_tokenizer_missing_merges(tmp_path):
 
     with pytest.raises(FileNotFoundError):
         hash_tokenizer_config(tokenizer_path)
+
+
+
+
+from openverifiablellm.tokenizer.tokenize_dataset import tokenize_dataset
+
+
+class DummyTokenizer:
+    def encode(self, text):
+        return [ord(c) for c in text]
+
+
+def test_tokenize_dataset_creates_output(tmp_path):
+
+    dataset = tmp_path / "dataset.txt"
+    dataset.write_text("hello\nworld")
+
+    output = tmp_path / "tokens.bin"
+
+    tokenizer = DummyTokenizer()
+
+    total_tokens = tokenize_dataset(dataset, tokenizer, output)
+
+    assert output.exists()
+    assert total_tokens > 0
+
+
+def test_tokenize_dataset_deterministic(tmp_path):
+
+    dataset = tmp_path / "dataset.txt"
+    dataset.write_text("test data\nanother line")
+
+    output1 = tmp_path / "tokens1.bin"
+    output2 = tmp_path / "tokens2.bin"
+
+    tokenizer = DummyTokenizer()
+
+    tokenize_dataset(dataset, tokenizer, output1)
+    tokenize_dataset(dataset, tokenizer, output2)
+
+    data1 = np.fromfile(output1, dtype=np.uint32)
+    data2 = np.fromfile(output2, dtype=np.uint32)
+
+    assert np.array_equal(data1, data2)