support custom training data

nanochat/dataloader.py

@@ -6,8 +6,17 @@ from nanochat.common import get_dist_info
 from nanochat.dataset import parquets_iter_batched
 from nanochat.tokenizer import get_tokenizer
 
-def tokenizing_distributed_data_loader(B, T, split, tokenizer_threads=4, tokenizer_batch_size=128, device="cuda"):
-    """Stream pretraining text from parquet files, tokenize, yield training batches."""
+def tokenizing_distributed_data_loader(B, T, split, tokenizer_threads=4, tokenizer_batch_size=128, device="cuda", data_dir=None):
+    """Stream pretraining text from parquet files, tokenize, yield training batches.
+
+    Args:
+        B: batch size
+        T: sequence length
+        split: "train" or "val"
+        tokenizer_threads: number of threads for tokenization
+        tokenizer_batch_size: batch size for tokenization
+        data_dir: optional custom directory containing parquet files (None = use default)
+    """
     assert split in ["train", "val"], "split must be 'train' or 'val'"
     ddp, ddp_rank, ddp_local_rank, ddp_world_size = get_dist_info()
     needed_tokens = B * T + 1 # +1 is because we also need the target at the last token
@@ -21,7 +30,7 @@ def tokenizing_distributed_data_loader(B, T, split, tokenizer_threads=4, tokeniz
     def document_batches():
         while True:
             # batch will iterate in group size of the parquet files, usually e.g. 1024 rows
-            for batch in parquets_iter_batched(split=split, start=ddp_rank, step=ddp_world_size):
+            for batch in parquets_iter_batched(split=split, start=ddp_rank, step=ddp_world_size, data_dir=data_dir):
                 # for the tokenizer we might want to go in usually smaller batches, e.g. 128 rows
                 for i in range(0, len(batch), tokenizer_batch_size):
                     yield batch[i:i+tokenizer_batch_size]

nanochat/dataset.py

@@ -31,7 +31,11 @@ os.makedirs(DATA_DIR, exist_ok=True)
 # These functions are useful utilities to other modules, can/should be imported
 
 def list_parquet_files(data_dir=None):
-    """ Looks into a data dir and returns full paths to all parquet files. """
+    """Looks into a data dir and returns full paths to all parquet files.
+
+    Args:
+        data_dir: optional custom directory containing parquet files (None = use default DATA_DIR)
+    """
     data_dir = DATA_DIR if data_dir is None else data_dir
     parquet_files = sorted([
         f for f in os.listdir(data_dir)
@@ -40,14 +44,17 @@ def list_parquet_files(data_dir=None):
     parquet_paths = [os.path.join(data_dir, f) for f in parquet_files]
     return parquet_paths
 
-def parquets_iter_batched(split, start=0, step=1):
-    """
-    Iterate through the dataset, in batches of underlying row_groups for efficiency.
-    - split can be "train" or "val". the last parquet file will be val.
-    - start/step are useful for skipping rows in DDP. e.g. start=rank, step=world_size
+def parquets_iter_batched(split, start=0, step=1, data_dir=None):
+    """Iterate through the dataset, in batches of underlying row_groups for efficiency.
+
+    Args:
+        split: "train" or "val". the last parquet file will be val.
+        start: starting row group index (useful for DDP, e.g. start=rank)
+        step: step size for row groups (useful for DDP, e.g. step=world_size)
+        data_dir: optional custom directory containing parquet files (None = use default)
     """
     assert split in ["train", "val"], "split must be 'train' or 'val'"
-    parquet_paths = list_parquet_files()
+    parquet_paths = list_parquet_files(data_dir=data_dir)
     parquet_paths = parquet_paths[:-1] if split == "train" else parquet_paths[-1:]
     for filepath in parquet_paths:
         pf = pq.ParquetFile(filepath)

scripts/base_loss.py

@@ -22,6 +22,7 @@ split_tokens = 20*524288  # number of tokens to evaluate per split
 model_tag = None # optional model tag for the output directory name
 model_step = None # optional model step for the output directory name
 device_type = "" # cuda|cpu|mps (empty => autodetect)
+data_dir = "" # path to directory containing parquet files with 'text' column (empty string = use default)
 exec(open(os.path.join('nanochat', 'configurator.py')).read()) # overrides from command line or config file
 
 # Load the base model and the tokenizer
@@ -37,8 +38,9 @@ assert split_tokens % tokens_per_step == 0, "split_tokens must be divisible by t
 steps = split_tokens // tokens_per_step
 token_bytes = get_token_bytes(device=device)
 bpb_results = {}
+custom_data_dir = data_dir if data_dir else None
 for split_name in ["train", "val"]:
-    loader = tokenizing_distributed_data_loader(device_batch_size, sequence_len, split_name, device=device)
+    loader = tokenizing_distributed_data_loader(device_batch_size, sequence_len, split_name, device=device, data_dir=custom_data_dir)
     with autocast_ctx:
         bpb = evaluate_bpb(model, loader, steps, token_bytes)
     print0(f"{split_name} bpb: {bpb:.4f}")