Pretraining Data Preparation

README.md

@@ -96,6 +96,66 @@ You can also turn off a specific encoder after training has completed, for testi
 > all the embeddings of the previous tokens (on the word level). This is done since not all causal LMs support
 > cross-attention, and so we want to avoid using it, and rely on the self-attention mechanism instead.
 
+## Data Preparation
+
+You can prepare datasets offline using the `welt-prepare-data` CLI.
+It streams a HuggingFace dataset, samples raw text with unit-based limits, and writes sharded `.jsonl.gz` files:
+
+```shell
+welt-prepare-data \
+    --dataset_name HuggingFaceFW/fineweb \
+    --dataset_config sample-10BT \
+    --max_total_units 3200000000 \
+    --num_units_per_file 100000000 \
+    --max_seq_length 512 \
+    --seed 42 \
+    --output_path /scratch/data/pretrain
+```
+
+Multiple datasets can be prepared into the same output directory:
+
+```shell
+welt-prepare-data \
+    --dataset_name monology/pile-uncopyrighted \
+    --max_total_units 1000000000 \
+    --num_units_per_file 100000000 \
+    --max_seq_length 512 \
+    --output_path /scratch/data/pretrain
+```
+
+The output directory contains sharded `.jsonl.gz` files and a `metadata.json` per run:
+
+```
+/scratch/data/pretrain/
+├── fineweb-sample-10BT-00000.jsonl.gz
+├── fineweb-sample-10BT-00001.jsonl.gz
+├── pile-uncopyrighted-00000.jsonl.gz
+└── metadata.json
+```
+
+Then train using the preprocessed data:
+
+```shell
+welt-train config.yaml --preprocessed_data_path /scratch/data/pretrain
+```
+
+| Argument | Description |
+|----------|-------------|
+| `--dataset_name` | HuggingFace dataset identifier (required) |
+| `--dataset_config` | Dataset config name (optional) |
+| `--dataset_split` | Split to use (default: "train") |
+| `--text_column` | Column containing text (default: "text") |
+| `--text_template` | Python format string template (optional) |
+| `--language` | Language tag to store with each example (e.g., "eng_Latn") |
+| `--unit_type` | Unit type for counting: "words" or "chars" (default: "words") |
+| `--max_total_units` | Max total units to sample (optional) |
+| `--num_units_per_file` | Max units per shard file (optional) |
+| `--max_seq_length` | Max words per example; splits long documents using word segmentation |
+| `--max_bytes_per_word` | Max bytes per word for word segmentation (default: 126) |
+| `--seed` | Random seed for shuffling |
+| `--drop_remainder` | Drop partial chunks at document boundaries |
+| `--output_path` | Output directory path (required) |
+
 ## Training
 
 Training instructions are available in the [welt_training/README.md](./welt_training/README.md).

pyproject.toml

@@ -36,6 +36,7 @@ train = [
     "scikit-learn", # For "accuracy" metric in evaluate
     "sacrebleu", # For usual bleu/chrF metrics
     "wandb", # For experiment tracking
+    "zstandard", # For compressing the data
 ]
 
 dion = [

tests/test_prepare_data.py

@@ -0,0 +1,164 @@
+import glob
+import gzip
+import json
+import shutil
+import tempfile
+
+import pytest
+from datasets import load_dataset
+
+from welt_training.prepare_data import get_shard_prefix, main
+
+
+# --- get_shard_prefix ---
+
+
+def test_get_shard_prefix_with_org():
+    assert get_shard_prefix("HuggingFaceFW/fineweb", None) == "fineweb"
+
+
+def test_get_shard_prefix_with_config():
+    assert get_shard_prefix("HuggingFaceFW/fineweb", "sample-10BT") == "fineweb-sample-10BT"
+
+
+def test_get_shard_prefix_no_org():
+    assert get_shard_prefix("wikitext", "wikitext-2-raw-v1") == "wikitext-wikitext-2-raw-v1"
+
+
+# --- Integration tests ---
+
+
+@pytest.fixture
+def temp_output_dir():
+    temp_dir = tempfile.mkdtemp(prefix="test_prepare_data_")
+    yield temp_dir
+    shutil.rmtree(temp_dir, ignore_errors=True)
+
+
+def test_prepare_data_creates_shards(temp_output_dir, monkeypatch):
+    """Test that welt-prepare-data creates sharded .jsonl.gz files that can be loaded."""
+    monkeypatch.setattr(
+        "sys.argv",
+        [
+            "welt-prepare-data",
+            "--dataset_name", "wikitext",
+            "--dataset_config", "wikitext-2-raw-v1",
+            "--max_total_units", "500",
+            "--num_units_per_file", "200",
+            "--seed", "42",
+            "--output_path", temp_output_dir,
+        ],
+    )
+    main()
+
+    # Verify shards were created
+    shard_files = sorted(glob.glob(f"{temp_output_dir}/*.jsonl.gz"))
+    assert len(shard_files) >= 2, f"Expected at least 2 shards, got {len(shard_files)}"
+
+    # Verify metadata
+    with open(f"{temp_output_dir}/metadata.json") as f:
+        metadata = json.load(f)
+    assert metadata["format"] == "welt-preprocessed-v1"
+    assert metadata["total_units"] <= 500
+    assert metadata["unit_type"] == "words"
+    assert metadata["num_shards"] == len(shard_files)
+
+    # Verify each shard is valid gzipped JSONL with a "text" field
+    total_examples = 0
+    for path in shard_files:
+        with gzip.open(path, "rt") as f:
+            for line in f:
+                example = json.loads(line)
+                assert "text" in example
+                assert isinstance(example["text"], str)
+                total_examples += 1
+    assert total_examples == metadata["num_examples"]
+
+    # Verify loading with HuggingFace datasets (same path as train.py)
+    ds = load_dataset("json", data_files=shard_files, split="train")
+    assert len(ds) == total_examples
+    assert "text" in ds.features
+
+
+def test_prepare_data_with_language(temp_output_dir, monkeypatch):
+    """Test that --language stores language metadata in each example."""
+    monkeypatch.setattr(
+        "sys.argv",
+        [
+            "welt-prepare-data",
+            "--dataset_name", "wikitext",
+            "--dataset_config", "wikitext-2-raw-v1",
+            "--max_total_units", "200",
+            "--language", "eng_Latn",
+            "--seed", "42",
+            "--output_path", temp_output_dir,
+        ],
+    )
+    main()
+
+    shard_files = sorted(glob.glob(f"{temp_output_dir}/*.jsonl.gz"))
+    for path in shard_files:
+        with gzip.open(path, "rt") as f:
+            for line in f:
+                example = json.loads(line)
+                assert example["language"] == "eng_Latn"
+
+    with open(f"{temp_output_dir}/metadata.json") as f:
+        metadata = json.load(f)
+    assert metadata["language"] == "eng_Latn"
+
+
+def test_prepare_data_unit_type_chars(temp_output_dir, monkeypatch):
+    """Test that --unit_type chars counts characters instead of words."""
+    monkeypatch.setattr(
+        "sys.argv",
+        [
+            "welt-prepare-data",
+            "--dataset_name", "wikitext",
+            "--dataset_config", "wikitext-2-raw-v1",
+            "--max_total_units", "500",
+            "--unit_type", "chars",
+            "--seed", "42",
+            "--output_path", temp_output_dir,
+        ],
+    )
+    main()
+
+    with open(f"{temp_output_dir}/metadata.json") as f:
+        metadata = json.load(f)
+    assert metadata["unit_type"] == "chars"
+    assert metadata["total_units"] <= 500
+
+
+def test_prepare_data_with_max_seq_length(temp_output_dir, monkeypatch):
+    """Test that --max_seq_length splits long documents into chunks."""
+    monkeypatch.setattr(
+        "sys.argv",
+        [
+            "welt-prepare-data",
+            "--dataset_name", "wikitext",
+            "--dataset_config", "wikitext-2-raw-v1",
+            "--max_total_units", "500",
+            "--max_seq_length", "32",
+            "--seed", "42",
+            "--output_path", temp_output_dir,
+        ],
+    )
+    main()
+
+    with open(f"{temp_output_dir}/metadata.json") as f:
+        metadata = json.load(f)
+    assert metadata["max_seq_length"] == 32
+    assert metadata["total_units"] <= 500
+
+    # Verify each example has at most max_seq_length words
+    from words_segmentation.tokenizer import WordsSegmentationTokenizer
+    pretokenizer = WordsSegmentationTokenizer()
+
+    shard_files = sorted(glob.glob(f"{temp_output_dir}/*.jsonl.gz"))
+    for path in shard_files:
+        with gzip.open(path, "rt") as f:
+            for line in f:
+                example = json.loads(line)
+                words = pretokenizer.tokenize(example["text"])
+                assert len(words) <= 32

welt_training/args_data.py

@@ -96,6 +96,10 @@ class DataTrainingArguments:
     keep_linebreaks: bool = field(
         default=True, metadata={"help": "Whether to keep line breaks when using TXT files or not."}
     )
+    preprocessed_data_path: str | None = field(
+        default=None,
+        metadata={"help": "Path to preprocessed dataset (from welt-prepare-data). Skips download and pretokenization."},
+    )
 
     def __post_init__(self):
         if self.streaming:
@@ -117,8 +121,13 @@ def __post_init__(self):
                 )
                 raise ValueError(msg)
 
-        if self.dataset_name is None and self.train_file is None and self.validation_file is None:
-            raise ValueError("Need either a dataset name or a training/validation file.")
+        if (
+            self.dataset_name is None
+            and self.train_file is None
+            and self.validation_file is None
+            and self.preprocessed_data_path is None
+        ):
+            raise ValueError("Need either a dataset name, a training/validation file, or a preprocessed data path.")
         else:
             if self.train_file is not None:
                 extension = self.train_file.split(".")[-1]