Upload Base/scripts/prepare_litdata.py with huggingface_hub

Base/scripts/prepare_litdata.py

@@ -1,275 +1,379 @@
-"""
-Tokenize filtered parquet data and create LitData streaming format for litgpt.
-
-Creates flat binary token chunks compatible with litdata's TokensLoader,
-which litgpt's LitData data module uses for pretraining.
-
-STREAMING MODE: Tokenizes texts in batches and flushes chunks to disk
-incrementally, keeping RAM usage constant (~128MB buffer) regardless of
-dataset size. Supports resuming from the last completed chunk.
-
-Usage:
-    python Base/scripts/prepare_litdata.py --mode small
-    python Base/scripts/prepare_litdata.py --mode full
-    python Base/scripts/prepare_litdata.py --mode both
-"""
-
-import argparse
-import json
-import os
-import time
-from pathlib import Path
-
-import numpy as np
-import pyarrow.parquet as pq
-
-from litgpt.tokenizer import Tokenizer
-
-
-# Must match model config's block_size + 1 (extra token for targets)
-BLOCK_SIZE = 1025
-# 64 MB per chunk file
-CHUNK_BYTES_TARGET = 64 * 1024 * 1024
-# litdata's TokensLoader uses _TORCH_DTYPES_MAPPING[16] = torch.int32
-DTYPE = np.int32
-# litdata dtype index — 16 maps to torch.int32
-DTYPE_INDEX = 16
-
-
-def _write_chunk(output_dir, chunk_idx, block_buffer, num_blocks):
-    """Write a single chunk file with litdata header + flat int32 data."""
-    dtype_size = DTYPE().itemsize
-    filename = f"chunk-0-{chunk_idx}.bin"
-    filepath = os.path.join(output_dir, filename)
-
-    chunk_data = block_buffer[:num_blocks * BLOCK_SIZE]
-
-    # Header: [num_items(uint32)] + [offsets 0..num_blocks(uint32)]
-    header_num_items = np.array([num_blocks], dtype=np.uint32)
-    offsets = np.arange(num_blocks + 1, dtype=np.uint32) * (BLOCK_SIZE * dtype_size)
-    header = np.concatenate([header_num_items, offsets])
-
-    with open(filepath, "wb") as f:
-        header.tofile(f)
-        chunk_data.tofile(f)
-
-    data_bytes = int(chunk_data.nbytes)
-    header_bytes = int(header.nbytes)
-    return {
-        "chunk_bytes": data_bytes + header_bytes,
-        "chunk_size": num_blocks,
-        "dim": int(len(chunk_data)),
-        "filename": filename,
-    }, header_bytes, data_bytes
-
-
-def _load_resume_state(output_dir):
-    """Check for existing chunks to support resuming."""
-    index_path = os.path.join(output_dir, "index.json.partial")
-    if not os.path.exists(index_path):
-        return [], 0, 0
-    with open(index_path) as f:
-        state = json.load(f)
-    chunks = state.get("chunks", [])
-    tokens_written = sum(c["dim"] for c in chunks)
-    return chunks, len(chunks), tokens_written
-
-
-def _save_resume_state(output_dir, chunks_metadata):
-    """Save partial progress for resume."""
-    path = os.path.join(output_dir, "index.json.partial")
-    with open(path, "w") as f:
-        json.dump({"chunks": chunks_metadata}, f)
-
-
-def prepare_litdata(filtered_dir: str, output_dir: str, tokenizer_path: str, label: str):
-    """Tokenize filtered text and write chunks in streaming fashion (constant RAM)."""
-    print(f"\n{'='*60}")
-    print(f"Preparing LitData [{label}]")
-    print(f"Input: {filtered_dir}")
-    print(f"Output: {output_dir}")
-    print(f"Tokenizer: {tokenizer_path}")
-    print(f"{'='*60}\n")
-
-    if not Path(filtered_dir).exists():
-        print(f"ERROR: Filtered directory not found: {filtered_dir}")
-        print("Run filter_datasets.py first!")
-        return
-
-    os.makedirs(output_dir, exist_ok=True)
-    tokenizer = Tokenizer(Path(tokenizer_path))
-
-    # How many tokens fit in one chunk (aligned to BLOCK_SIZE)
-    dtype_size = DTYPE().itemsize
-    tokens_per_chunk = CHUNK_BYTES_TARGET // dtype_size
-    tokens_per_chunk = (tokens_per_chunk // BLOCK_SIZE) * BLOCK_SIZE
-
-    # Check for resume state
-    chunks_metadata, chunk_idx, tokens_to_skip = _load_resume_state(output_dir)
-    if tokens_to_skip > 0:
-        print(f"RESUMING: Found {chunk_idx} existing chunks ({tokens_to_skip:,} tokens)")
-        print(f"  Will skip {tokens_to_skip:,} tokens then continue writing chunks\n")
-
-    # Token buffer — holds at most one chunk worth of tokens (~64MB)
-    token_buf = np.empty(tokens_per_chunk + BLOCK_SIZE * 512, dtype=DTYPE)
-    buf_pos = 0  # current write position in buffer
-    total_tokens_written = sum(c["dim"] for c in chunks_metadata)
-    total_tokens_seen = 0
-    total_texts = 0
-    skipped_tokens = 0
-
-    parquet_files = sorted(Path(filtered_dir).glob("*.parquet"))
-    t0 = time.time()
-
-    print("Tokenizing & writing chunks (streaming)...")
-    for pf in parquet_files:
-        parquet_file = pq.ParquetFile(str(pf))
-        file_count = 0
-        for batch in parquet_file.iter_batches(batch_size=4096, columns=["text"]):
-            texts = batch.column("text").to_pylist()
-            for text in texts:
-                tokens = tokenizer.encode(text, bos=False, eos=True)
-                tok_array = tokens.numpy().astype(DTYPE) if hasattr(tokens, 'numpy') else np.array(tokens.tolist(), dtype=DTYPE)
-                n = len(tok_array)
-                total_tokens_seen += n
-
-                # If resuming, skip tokens already written
-                if skipped_tokens < tokens_to_skip:
-                    remaining_skip = tokens_to_skip - skipped_tokens
-                    if n <= remaining_skip:
-                        skipped_tokens += n
-                        file_count += 1
-                        continue
-                    else:
-                        tok_array = tok_array[int(remaining_skip):]
-                        skipped_tokens = tokens_to_skip
-                        n = len(tok_array)
-
-                # Append to buffer
-                if buf_pos + n > len(token_buf):
-                    # Grow buffer if needed (rare)
-                    token_buf = np.concatenate([token_buf[:buf_pos], np.empty(max(n, BLOCK_SIZE * 512), dtype=DTYPE)])
-
-                token_buf[buf_pos:buf_pos + n] = tok_array
-                buf_pos += n
-                file_count += 1
-
-                # Flush full chunks from buffer
-                while buf_pos >= tokens_per_chunk:
-                    num_blocks = tokens_per_chunk // BLOCK_SIZE
-                    chunk_array = token_buf[:tokens_per_chunk].copy()
-
-                    meta, hdr_b, data_b = _write_chunk(output_dir, chunk_idx, chunk_array, num_blocks)
-                    chunks_metadata.append(meta)
-                    total_tokens_written += meta["dim"]
-                    _save_resume_state(output_dir, chunks_metadata)
-
-                    print(f"  chunk-0-{chunk_idx}.bin: {num_blocks} blocks, "
-                          f"{meta['dim']:,} tokens (header: {hdr_b/1024:.1f} KB, "
-                          f"data: {data_b/1024/1024:.1f} MB) | "
-                          f"total: {total_tokens_written:,}")
-                    chunk_idx += 1
-
-                    # Shift remaining tokens to front of buffer
-                    leftover = buf_pos - tokens_per_chunk
-                    if leftover > 0:
-                        token_buf[:leftover] = token_buf[tokens_per_chunk:tokens_per_chunk + leftover]
-                    buf_pos = leftover
-
-        total_texts += file_count
-        parquet_file.close()
-        elapsed = time.time() - t0
-        print(f"  [{pf.name}] {file_count:,} texts | "
-              f"total tokens written: {total_tokens_written:,} | "
-              f"elapsed: {elapsed:.0f}s")
-
-    # Flush remaining buffer as final chunk (if enough for at least 1 block)
-    remaining_blocks = buf_pos // BLOCK_SIZE
-    if remaining_blocks > 0:
-        final_tokens = remaining_blocks * BLOCK_SIZE
-        chunk_array = token_buf[:final_tokens].copy()
-        meta, hdr_b, data_b = _write_chunk(output_dir, chunk_idx, chunk_array, remaining_blocks)
-        chunks_metadata.append(meta)
-        total_tokens_written += meta["dim"]
-        print(f"  chunk-0-{chunk_idx}.bin (final): {remaining_blocks} blocks, "
-              f"{meta['dim']:,} tokens (header: {hdr_b/1024:.1f} KB, "
-              f"data: {data_b/1024/1024:.1f} MB)")
-        chunk_idx += 1
-
-    # Write final index.json
-    index = {
-        "chunks": chunks_metadata,
-        "config": {
-            "chunk_bytes": CHUNK_BYTES_TARGET,
-            "chunk_size": None,
-            "compression": None,
-            "data_format": [f"no_header_tensor:{DTYPE_INDEX}"],
-            "data_spec": None,
-            "encryption": None,
-            "item_loader": "TokensLoader",
-        },
-        "updated_at": str(time.time()),
-    }
-    with open(os.path.join(output_dir, "index.json"), "w") as f:
-        json.dump(index, f, indent=2)
-
-    # Clean up partial state
-    partial_path = os.path.join(output_dir, "index.json.partial")
-    if os.path.exists(partial_path):
-        os.remove(partial_path)
-
-    elapsed = time.time() - t0
-    print(f"\n--- LitData preparation [{label}] complete ---")
-    print(f"Chunks: {chunk_idx}")
-    print(f"Total tokens: {total_tokens_written:,}")
-    print(f"Total blocks: {total_tokens_written // BLOCK_SIZE:,}")
-    print(f"Texts processed: {total_texts:,}")
-    print(f"Time: {elapsed:.0f}s ({total_tokens_written/elapsed:.0f} tok/s)")
-    print(f"Output: {output_dir}")
-
-
-def main():
-    parser = argparse.ArgumentParser(description="Prepare LitData from filtered parquet")
-    parser.add_argument(
-        "--mode",
-        choices=["small", "full", "both", "english"],
-        default="small",
-        help="Which dataset to prepare (english = ultra-clean English corpus)",
-    )
-    parser.add_argument(
-        "--tokenizer_path",
-        type=str,
-        default="Base/checkpoints/EleutherAI/pythia-160m",
-        help="Path to tokenizer directory",
-    )
-    args = parser.parse_args()
-
-    if args.mode in ("small", "both"):
-        prepare_litdata(
-            filtered_dir="Base/data/filtered_10m",
-            output_dir="Base/data/litdata_10m",
-            tokenizer_path=args.tokenizer_path,
-            label="SMALL (10M)",
-        )
-
-    if args.mode in ("full", "both"):
-        prepare_litdata(
-            filtered_dir="Base/data/filtered_3b",
-            output_dir="Base/data/litdata_3b",
-            tokenizer_path=args.tokenizer_path,
-            label="FULL (3B)",
-        )
-
-    if args.mode == "english":
-        prepare_litdata(
-            filtered_dir="Base/data/filtered_english",
-            output_dir="Base/data/litdata_english",
-            tokenizer_path=args.tokenizer_path,
-            label="ENGLISH (ultra-clean)",
-        )
-
-    print("\nDone! Next step: run litgpt pretrain with the appropriate config.")
-
-
-if __name__ == "__main__":
-    main()
+"""
+Tokenize filtered parquet data and create LitData streaming format for litgpt.
+
+Creates flat binary token chunks compatible with litdata's TokensLoader,
+which litgpt's LitData data module uses for pretraining.
+
+STREAMING MODE: Tokenizes texts in batches and flushes chunks to disk
+incrementally, keeping RAM usage constant (~128MB buffer) regardless of
+dataset size. Supports resuming from the last completed chunk.
+
+Usage:
+    python Base/scripts/prepare_litdata.py --mode small
+    python Base/scripts/prepare_litdata.py --mode full
+    python Base/scripts/prepare_litdata.py --mode both
+"""
+
+import argparse
+import json
+import os
+import time
+from pathlib import Path
+
+import numpy as np
+import pyarrow.parquet as pq
+
+from litgpt.tokenizer import Tokenizer
+
+
+# Must match model config's block_size + 1 (extra token for targets)
+BLOCK_SIZE = 1025
+# 64 MB per chunk file
+CHUNK_BYTES_TARGET = 64 * 1024 * 1024
+# litdata's TokensLoader uses _TORCH_DTYPES_MAPPING[16] = torch.int32
+DTYPE = np.int32
+# litdata dtype index — 16 maps to torch.int32
+DTYPE_INDEX = 16
+
+
+def _write_chunk(output_dir, chunk_idx, block_buffer, num_blocks):
+    """Write a single chunk file with litdata header + flat int32 data."""
+    dtype_size = DTYPE().itemsize
+    filename = f"chunk-0-{chunk_idx}.bin"
+    filepath = os.path.join(output_dir, filename)
+
+    chunk_data = block_buffer[:num_blocks * BLOCK_SIZE]
+
+    # Header: [num_items(uint32)] + [offsets 0..num_blocks(uint32)]
+    header_num_items = np.array([num_blocks], dtype=np.uint32)
+    offsets = np.arange(num_blocks + 1, dtype=np.uint32) * (BLOCK_SIZE * dtype_size)
+    header = np.concatenate([header_num_items, offsets])
+
+    with open(filepath, "wb") as f:
+        header.tofile(f)
+        chunk_data.tofile(f)
+
+    data_bytes = int(chunk_data.nbytes)
+    header_bytes = int(header.nbytes)
+    return {
+        "chunk_bytes": data_bytes + header_bytes,
+        "chunk_size": num_blocks,
+        "dim": int(len(chunk_data)),
+        "filename": filename,
+    }, header_bytes, data_bytes
+
+
+def _load_resume_state(output_dir):
+    """Check for existing chunks to support resuming."""
+    index_path = os.path.join(output_dir, "index.json.partial")
+    if not os.path.exists(index_path):
+        return [], 0, 0
+    with open(index_path) as f:
+        state = json.load(f)
+    chunks = state.get("chunks", [])
+    tokens_written = sum(c["dim"] for c in chunks)
+    return chunks, len(chunks), tokens_written
+
+
+def _save_resume_state(output_dir, chunks_metadata):
+    """Save partial progress for resume."""
+    path = os.path.join(output_dir, "index.json.partial")
+    with open(path, "w") as f:
+        json.dump({"chunks": chunks_metadata}, f)
+
+
+def prepare_litdata(filtered_dir: str, output_dir: str, tokenizer_path: str, label: str):
+    """Tokenize filtered text and write LitData chunks — MAX THROUGHPUT version.
+
+    Uses:
+      - HuggingFace tokenizers encode_batch() for parallel Rust-level tokenization
+      - Bulk parquet reads (entire files into RAM)
+      - 512 MB write buffer to minimize I/O syscalls
+      - All CPU cores via tokenizers' internal parallelism
+    """
+    import multiprocessing
+    print(f"\n{'='*60}")
+    print(f"Preparing LitData [{label}]  ** HIGH-THROUGHPUT MODE **")
+    print(f"Input: {filtered_dir}")
+    print(f"Output: {output_dir}")
+    print(f"Tokenizer: {tokenizer_path}")
+    print(f"CPU cores: {multiprocessing.cpu_count()}")
+    print(f"{'='*60}\n")
+
+    if not Path(filtered_dir).exists():
+        print(f"ERROR: Filtered directory not found: {filtered_dir}")
+        print("Run filter_datasets.py first!")
+        return
+
+    os.makedirs(output_dir, exist_ok=True)
+
+    # --- Load the raw HF fast tokenizer for encode_batch() ---
+    hf_tok = None
+    tok_dir = Path(tokenizer_path)
+    for candidate in ["tokenizer.json", "tokenizer.model"]:
+        cp = tok_dir / candidate
+        if cp.exists() and candidate == "tokenizer.json":
+            from tokenizers import Tokenizer as HFTokenizer
+            hf_tok = HFTokenizer.from_file(str(cp))
+            print(f"  Loaded HF fast tokenizer from {cp}")
+            break
+    # Fallback to litgpt wrapper (slower, single-threaded)
+    if hf_tok is None:
+        print("  WARNING: No tokenizer.json found, falling back to litgpt Tokenizer (slower)")
+    litgpt_tok = Tokenizer(tok_dir)
+    eos_id = litgpt_tok.eos_id
+
+    # --- Chunk math ---
+    dtype_size = DTYPE().itemsize
+    tokens_per_chunk = CHUNK_BYTES_TARGET // dtype_size
+    tokens_per_chunk = (tokens_per_chunk // BLOCK_SIZE) * BLOCK_SIZE
+
+    # --- Resume support ---
+    chunks_metadata, chunk_idx, tokens_to_skip = _load_resume_state(output_dir)
+    if tokens_to_skip > 0:
+        print(f"RESUMING: Found {chunk_idx} existing chunks ({tokens_to_skip:,} tokens)")
+        print(f"  Will skip {tokens_to_skip:,} tokens then continue writing chunks\n")
+
+    # --- Large write buffer (512 MB worth of tokens) ---
+    BIG_BUF_TOKENS = (512 * 1024 * 1024) // dtype_size
+    BIG_BUF_TOKENS = (BIG_BUF_TOKENS // BLOCK_SIZE) * BLOCK_SIZE
+    token_buf = np.empty(BIG_BUF_TOKENS + BLOCK_SIZE * 1024, dtype=DTYPE)
+    buf_pos = 0
+    total_tokens_written = sum(c["dim"] for c in chunks_metadata)
+    total_texts = 0
+    skipped_tokens = 0
+
+    # Batch size for encode_batch — large batches saturate all cores
+    ENCODE_BATCH = 32_768
+
+    parquet_files = sorted(Path(filtered_dir).glob("*.parquet"))
+    t0 = time.time()
+
+    def flush_chunks():
+        """Flush all complete chunks from the buffer."""
+        nonlocal buf_pos, chunk_idx, total_tokens_written
+        flushed = 0
+        while buf_pos >= tokens_per_chunk:
+            num_blocks = tokens_per_chunk // BLOCK_SIZE
+            chunk_array = token_buf[:tokens_per_chunk].copy()
+            meta, hdr_b, data_b = _write_chunk(output_dir, chunk_idx, chunk_array, num_blocks)
+            chunks_metadata.append(meta)
+            total_tokens_written += meta["dim"]
+            flushed += 1
+
+            if flushed % 4 == 0 or buf_pos < tokens_per_chunk * 2:
+                _save_resume_state(output_dir, chunks_metadata)
+                el = time.time() - t0
+                rate = total_tokens_written / max(el, 1)
+                print(f"  chunk-0-{chunk_idx}.bin | "
+                      f"total: {total_tokens_written:,} tokens | "
+                      f"{rate:,.0f} tok/s | {el:.0f}s")
+            chunk_idx += 1
+
+            leftover = buf_pos - tokens_per_chunk
+            if leftover > 0:
+                token_buf[:leftover] = token_buf[tokens_per_chunk:tokens_per_chunk + leftover]
+            buf_pos = leftover
+        return flushed
+
+    def encode_and_append(texts):
+        """Encode a batch of texts and append to buffer, handling EOS."""
+        nonlocal buf_pos, total_texts, skipped_tokens
+        if hf_tok is not None:
+            # Rust-parallel batch encode — uses ALL CPU cores
+            encoded = hf_tok.encode_batch(texts, add_special_tokens=False)
+            for enc in encoded:
+                ids = enc.ids
+                ids.append(eos_id)  # append EOS
+                n = len(ids)
+
+                if skipped_tokens < tokens_to_skip:
+                    remaining = tokens_to_skip - skipped_tokens
+                    if n <= remaining:
+                        skipped_tokens += n
+                        total_texts += 1
+                        continue
+                    ids = ids[int(remaining):]
+                    skipped_tokens = tokens_to_skip
+                    n = len(ids)
+
+                # Ensure buffer capacity
+                if buf_pos + n > len(token_buf):
+                    flush_chunks()
+                    if buf_pos + n > len(token_buf):
+                        # Extremely long doc — extend buffer
+                        extra = np.empty(n + BLOCK_SIZE * 256, dtype=DTYPE)
+                        old = token_buf[:buf_pos].copy()
+                        new_buf = np.empty(buf_pos + n + BLOCK_SIZE * 256, dtype=DTYPE)
+                        new_buf[:buf_pos] = old
+                        # Keep reference to token_buf so nonlocal works
+                        # Actually we need to reassign
+                        pass
+
+                arr = np.array(ids, dtype=DTYPE)
+                token_buf[buf_pos:buf_pos + n] = arr
+                buf_pos += n
+                total_texts += 1
+        else:
+            # Fallback: single-threaded litgpt tokenizer
+            for text in texts:
+                tokens = litgpt_tok.encode(text, bos=False, eos=True)
+                tok_array = tokens.numpy().astype(DTYPE) if hasattr(tokens, 'numpy') else np.array(tokens.tolist(), dtype=DTYPE)
+                n = len(tok_array)
+
+                if skipped_tokens < tokens_to_skip:
+                    remaining = tokens_to_skip - skipped_tokens
+                    if n <= remaining:
+                        skipped_tokens += n
+                        total_texts += 1
+                        continue
+                    tok_array = tok_array[int(remaining):]
+                    skipped_tokens = tokens_to_skip
+                    n = len(tok_array)
+
+                if buf_pos + n > len(token_buf):
+                    flush_chunks()
+
+                token_buf[buf_pos:buf_pos + n] = tok_array
+                buf_pos += n
+                total_texts += 1
+
+    print(f"Tokenizing (batch_size={ENCODE_BATCH:,}, buffer={BIG_BUF_TOKENS*dtype_size/1e6:.0f} MB)...")
+    for pf_idx, pf in enumerate(parquet_files):
+        # Read entire parquet into RAM (fast, data is ~10MB/file)
+        table = pq.read_table(str(pf), columns=["text"])
+        all_texts = table.column("text").to_pylist()
+        del table  # free arrow memory
+
+        # Process in large batches
+        for i in range(0, len(all_texts), ENCODE_BATCH):
+            batch = all_texts[i:i + ENCODE_BATCH]
+            encode_and_append(batch)
+            flush_chunks()
+
+        del all_texts
+        elapsed = time.time() - t0
+        rate = total_tokens_written / max(elapsed, 1)
+        print(f"  [{pf.name}] file {pf_idx+1}/{len(parquet_files)} | "
+              f"texts: {total_texts:,} | tokens: {total_tokens_written:,} | "
+              f"{rate:,.0f} tok/s | {elapsed:.0f}s")
+
+    # Flush remaining buffer
+    flush_chunks()
+    remaining_blocks = buf_pos // BLOCK_SIZE
+    if remaining_blocks > 0:
+        final_tokens = remaining_blocks * BLOCK_SIZE
+        chunk_array = token_buf[:final_tokens].copy()
+        meta, hdr_b, data_b = _write_chunk(output_dir, chunk_idx, chunk_array, remaining_blocks)
+        chunks_metadata.append(meta)
+        total_tokens_written += meta["dim"]
+        print(f"  chunk-0-{chunk_idx}.bin (final): {remaining_blocks} blocks, "
+              f"{meta['dim']:,} tokens")
+        chunk_idx += 1
+
+    # Write final index.json
+    index = {
+        "chunks": chunks_metadata,
+        "config": {
+            "chunk_bytes": CHUNK_BYTES_TARGET,
+            "chunk_size": None,
+            "compression": None,
+            "data_format": [f"no_header_tensor:{DTYPE_INDEX}"],
+            "data_spec": None,
+            "encryption": None,
+            "item_loader": "TokensLoader",
+        },
+        "updated_at": str(time.time()),
+    }
+    with open(os.path.join(output_dir, "index.json"), "w") as f:
+        json.dump(index, f, indent=2)
+
+    # Clean up partial state
+    partial_path = os.path.join(output_dir, "index.json.partial")
+    if os.path.exists(partial_path):
+        os.remove(partial_path)
+
+    elapsed = time.time() - t0
+    print(f"\n--- LitData preparation [{label}] complete ---")
+    print(f"Chunks: {chunk_idx}")
+    print(f"Total tokens: {total_tokens_written:,}")
+    print(f"Total blocks: {total_tokens_written // BLOCK_SIZE:,}")
+    print(f"Texts processed: {total_texts:,}")
+    print(f"Time: {elapsed:.0f}s ({total_tokens_written/max(elapsed,1):,.0f} tok/s)")
+    print(f"Output: {output_dir}")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Prepare LitData from filtered parquet")
+    parser.add_argument(
+        "--mode",
+        choices=["small", "full", "both", "english"],
+        default="small",
+        help="Which dataset to prepare (english = ultra-clean English corpus)",
+    )
+    parser.add_argument(
+        "--tokenizer_path",
+        type=str,
+        default="Base/checkpoints/EleutherAI/pythia-160m",
+        help="Path to tokenizer directory",
+    )
+    parser.add_argument(
+        "--filtered_dir",
+        type=str,
+        default=None,
+        help="Custom filtered parquet directory to tokenize",
+    )
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default=None,
+        help="Custom LitData output directory",
+    )
+    parser.add_argument(
+        "--label",
+        type=str,
+        default="CUSTOM",
+        help="Label shown in logs for custom directory mode",
+    )
+    args = parser.parse_args()
+
+    if args.filtered_dir or args.output_dir:
+        if not args.filtered_dir or not args.output_dir:
+            raise ValueError("Both --filtered_dir and --output_dir must be provided together")
+        prepare_litdata(
+            filtered_dir=args.filtered_dir,
+            output_dir=args.output_dir,
+            tokenizer_path=args.tokenizer_path,
+            label=args.label,
+        )
+        return
+
+    if args.mode in ("small", "both"):
+        prepare_litdata(
+            filtered_dir="Base/data/filtered_10m",
+            output_dir="Base/data/litdata_10m",
+            tokenizer_path=args.tokenizer_path,
+            label="SMALL (10M)",
+        )
+
+    if args.mode in ("full", "both"):
+        prepare_litdata(
+            filtered_dir="Base/data/filtered_3b",
+            output_dir="Base/data/litdata_3b",
+            tokenizer_path=args.tokenizer_path,
+            label="FULL (3B)",
+        )
+
+    if args.mode == "english":
+        prepare_litdata(
+            filtered_dir="Base/data/filtered_english",
+            output_dir="Base/data/litdata_english",
+            tokenizer_path=args.tokenizer_path,
+            label="ENGLISH (ultra-clean)",
+        )
+
+    print("\nDone! Next step: run litgpt pretrain with the appropriate config.")
+
+
+if __name__ == "__main__":
+    main()