feat: implement dataset preparation and streaming

This commit implements the dataset preparation phase of the roadmap.
It includes a script to download and tokenize datasets (FineWeb-Edu and GSM8K) into efficient binary format, and a new Dataset class to stream this data during training.

This fulfills the requirement to design and prepare a dataset for training.
It provides:
1. `scripts/data/prepare_dataset.py`: A script to download Hugging Face datasets (FineWeb-Edu, GSM8K), tokenize them, and save them as efficient `uint16` binary files (saving storage).
2. `crsm/dataset.py`: A new `PretokenizedDataset` class that streams these binary files using `numpy.memmap` (saving memory) and correctly stitches data across file boundaries to prevent token loss during streaming.

Author: Pomilon

crsm/dataset.py

@@ -3,12 +3,12 @@
 """
 import gc
 from collections import deque
-from typing import Iterator, Optional
+from typing import Iterator, Optional, List, Tuple, Iterable
 import torch
 from torch.utils.data import Dataset, IterableDataset
 from pathlib import Path
 from collections import Counter
-from typing import List, Tuple, Iterable
+import numpy as np
 
 from .tokenizer import Tokenizer
 
@@ -45,6 +45,132 @@ def encode(self, tokens: List[str]) -> List[int]:
         return [self.stoi.get(t, self.stoi['<unk>']) for t in tokens]
 
 
+class PretokenizedDataset(IterableDataset):
+    """
+    Dataset that streams tokens from pre-tokenized binary files.
+    """
+    def __init__(self, data_dir: str, seq_len: int = 1024, split: str = "train", 
+                 vocab_size: int = 50257):
+        self.data_dir = Path(data_dir)
+        self.seq_len = seq_len
+        self.vocab_size = vocab_size
+        
+        # Check vocab size safety
+        if self.vocab_size > 65535:
+            raise ValueError(f"Vocab size {self.vocab_size} too large for uint16 storage used by PretokenizedDataset")
+        
+        # Find all bin files matching the split
+        self.files = sorted(list(self.data_dir.glob(f"*{split}*.bin")))
+        if not self.files:
+             raise FileNotFoundError(f"No bin files found in {data_dir} for split '{split}'. Did you run prepare_dataset.py?")
+            
+        print(f"Found {len(self.files)} bin files in {data_dir} for split {split}")
+        
+    def __iter__(self):
+        worker_info = torch.utils.data.get_worker_info()
+        if worker_info:
+            # Shard files across workers
+            # Simple interleaving
+            files = self.files[worker_info.id::worker_info.num_workers]
+        else:
+            files = self.files
+            
+        # Buffer to hold tokens from multiple files to ensure continuity
+        buffer = np.array([], dtype=np.uint16)
+        chunk_len = self.seq_len + 1
+        
+        for path in files:
+            # Read file as numpy array
+            try:
+                # Use memmap for efficiency
+                new_data = np.memmap(path, dtype=np.uint16, mode='r')
+                
+                # Append to buffer (we have to copy here, but buffer is small)
+                # Actually, appending memmap to numpy array forces a read. 
+                # To be efficient, we should process the buffer as much as possible first.
+                
+                # However, for continuity, we need to stitch the end of the previous file 
+                # with the start of the new file.
+                
+                # Since we are iterating, we can yield chunks from `new_data` 
+                # but we need to handle the "leftover" from previous file.
+                
+                # If buffer is not empty (from previous iteration), prepend it
+                if len(buffer) > 0:
+                     # This forces a read of the whole file if we use np.concatenate
+                     # Ideally we only read the beginning.
+                     
+                     # Optimization: Only concat the leftovers with the start of new file?
+                     # No, memmap is array-like.
+                     
+                     # Let's just process the memmap directly and keep the tail in buffer.
+                     pass
+                
+                # We can't easily concatenate memmap without reading it. 
+                # Strategy: 
+                # 1. Take leftovers from buffer.
+                # 2. Iterate through memmap. 
+                # 3. If we have enough for a chunk using leftovers + start of memmap, yield it.
+                # 4. Then yield chunks from memmap.
+                # 5. Save tail of memmap to buffer.
+                
+                current_idx = 0
+                total_len = len(new_data)
+                
+                # Handle leftovers
+                if len(buffer) > 0:
+                    needed = chunk_len - len(buffer)
+                    if total_len >= needed:
+                        # Take 'needed' from new_data
+                        part = new_data[:needed] # This reads from disk
+                        full_chunk = np.concatenate([buffer, part])
+                        
+                        x = torch.from_numpy(full_chunk[:-1].astype(np.int64))
+                        y = torch.from_numpy(full_chunk[1:].astype(np.int64))
+                        yield x, y
+                        
+                        current_idx = needed
+                        buffer = np.array([], dtype=np.uint16)
+                    else:
+                        # File is too small to complete the buffer
+                        # Append all of it to buffer and continue to next file
+                        part = new_data[:]
+                        buffer = np.concatenate([buffer, part])
+                        continue
+                        
+                # Process main body of file
+                # We can compute how many chunks fit
+                remaining = total_len - current_idx
+                num_chunks = remaining // chunk_len
+                
+                if num_chunks > 0:
+                    # Create a view or slice
+                    # Note: slicing memmap returns memmap, which is good.
+                    # We can iterate over the slices.
+                    
+                    # But constructing individual tensors from memmap slices is fine.
+                    for i in range(num_chunks):
+                        start = current_idx + i * chunk_len
+                        end = start + chunk_len
+                        chunk = new_data[start:end]
+                        
+                        # Copy to memory and convert
+                        chunk_arr = np.array(chunk, dtype=np.int64)
+                        x = torch.from_numpy(chunk_arr[:-1])
+                        y = torch.from_numpy(chunk_arr[1:])
+                        yield x, y
+                        
+                    current_idx += num_chunks * chunk_len
+                    
+                # Save leftovers
+                if current_idx < total_len:
+                    buffer = np.array(new_data[current_idx:], dtype=np.uint16)
+                    
+            except Exception as e:
+                print(f"Error reading {path}: {e}")
+                continue
+
+
 class RealTextDataset(Dataset):
     """Loads text files from a directory, builds a simple token-level vocab,
     and yields sequences for next-token prediction.

notebooks/cloud_training/1_train_backbone.ipynb

@@ -36,7 +36,8 @@
    "source": [
     "!git clone https://github.com/Pomilon-Intelligence-Lab/crsm.git\n",
     "%cd crsm\n",
-    "!pip install -r requirements.txt"
+    "!pip install -r requirements.txt\n",
+    "!pip install -e ."
    ]
   },
   {
@@ -60,7 +61,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3",
+   "display_name": "Python 3 (ipykernel)",
    "language": "python",
    "name": "python3"
   },
@@ -74,7 +75,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.10"
+   "version": "3.12.12"
   }
  },
  "nbformat": 4,

notebooks/cloud_training/2_distill_dynamics.ipynb

@@ -36,7 +36,8 @@
    "source": [
     "!git clone https://github.com/Pomilon-Intelligence-Lab/crsm.git\n",
     "%cd crsm\n",
-    "!pip install -r requirements.txt"
+    "!pip install -r requirements.txt\n",
+    "!pip install -e ."
    ]
   },
   {

notebooks/cloud_training/3_train_judgment.ipynb

@@ -36,7 +36,8 @@
    "source": [
     "!git clone https://github.com/Pomilon-Intelligence-Lab/crsm.git\n",
     "%cd crsm\n",
-    "!pip install -r requirements.txt"
+    "!pip install -r requirements.txt\n",
+    "!pip install -e ."
    ]
   },
   {

scripts/data/prepare_dataset.py

@@ -0,0 +1,118 @@
+"""
+Prepare Dataset Script
+----------------------
+Downloads and tokenizes datasets (FineWeb-Edu, GSM8K) into efficient binary format (uint16).
+This allows for fast streaming during training with minimal memory overhead.
+
+Usage:
+    python scripts/data/prepare_dataset.py --dataset fineweb --output_dir data/fineweb
+    python scripts/data/prepare_dataset.py --dataset gsm8k --output_dir data/gsm8k
+"""
+import os
+import sys
+import argparse
+import numpy as np
+from pathlib import Path
+from tqdm import tqdm
+from datasets import load_dataset
+from transformers import AutoTokenizer
+import multiprocessing as mp
+
+# Ensure the package is in the path
+sys.path.insert(0, '.')
+
+def get_tokenizer(model_name="gpt2"):
+    return AutoTokenizer.from_pretrained(model_name)
+
+def process_fineweb(example, tokenizer):
+    ids = tokenizer.encode(example['text'])
+    ids.append(tokenizer.eos_token_id)
+    return ids
+
+def process_gsm8k(example, tokenizer):
+    # Format: Question: ... Answer: ...
+    text = f"Question: {example['question']}\nAnswer: {example['answer']}"
+    ids = tokenizer.encode(text)
+    ids.append(tokenizer.eos_token_id)
+    return ids
+
+def write_to_bin(tokens, output_file):
+    # Convert to uint16 (ensure vocab < 65535)
+    # GPT2 vocab is 50257, so it fits.
+    arr = np.array(tokens, dtype=np.uint16)
+    with open(output_file, "wb") as f:
+        f.write(arr.tobytes())
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--dataset', type=str, required=True, choices=['fineweb', 'gsm8k'], help="Dataset to prepare")
+    parser.add_argument('--output-dir', type=str, required=True, help="Directory to save binary files")
+    parser.add_argument('--tokenizer', type=str, default="gpt2", help="Tokenizer name")
+    parser.add_argument('--shard-size', type=int, default=100_000_000, help="Tokens per shard")
+    parser.add_argument('--subset', type=str, default="sample-10BT", help="Subset for FineWeb (default: sample-10BT)")
+    
+    args = parser.parse_args()
+    
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    
+    tokenizer = get_tokenizer(args.tokenizer)
+    print(f"Loaded tokenizer: {args.tokenizer} (vocab size: {tokenizer.vocab_size})")
+    
+    if tokenizer.vocab_size >= 65535:
+        print("Error: Tokenizer vocab size too large for uint16 storage!")
+        sys.exit(1)
+        
+    # Load Dataset
+    print(f"Loading dataset {args.dataset}...")
+    if args.dataset == 'fineweb':
+        # Use streaming for large datasets
+        ds = load_dataset("HuggingFaceFW/fineweb-edu", name=args.subset, split="train", streaming=True)
+        process_fn = process_fineweb
+        prefix = "fineweb"
+    elif args.dataset == 'gsm8k':
+        ds = load_dataset("gsm8k", "main", split="train", streaming=True)
+        process_fn = process_gsm8k
+        prefix = "gsm8k"
+        
+    token_buffer = []
+    shard_idx = 0
+    total_tokens = 0
+    
+    print(f"Processing and tokenizing to {output_dir}...")
+    
+    pbar = tqdm(desc="Tokens", unit="tok")
+    
+    for example in ds:
+        tokens = process_fn(example, tokenizer)
+        token_buffer.extend(tokens)
+        
+        # Update pbar occasionally
+        if len(token_buffer) % 1000 == 0:
+             pbar.update(len(tokens)) # Approximate update
+        
+        if len(token_buffer) >= args.shard_size:
+            # Write shard
+            filename = output_dir / f"{prefix}_train_{shard_idx:04d}.bin"
+            write_to_bin(token_buffer, filename)
+            
+            total_tokens += len(token_buffer)
+            print(f"Saved {filename} ({len(token_buffer)} tokens)")
+            
+            token_buffer = []
+            shard_idx += 1
+            
+            # For demonstration/testing purposes, we might stop early if running in CI/Test env
+            # But this is a script, so let it run.
+            
+    # Write remaining
+    if token_buffer:
+        filename = output_dir / f"{prefix}_train_{shard_idx:04d}.bin"
+        write_to_bin(token_buffer, filename)
+        total_tokens += len(token_buffer)
+        print(f"Saved {filename} ({len(token_buffer)} tokens)")
+        
+    print(f"Done. Total tokens: {total_tokens}")
+
+if __name__ == "__main__":
+    main()

tests/test_dataset_continuity.py

@@ -0,0 +1,78 @@
+
+import unittest
+import torch
+import numpy as np
+from pathlib import Path
+import shutil
+import tempfile
+from crsm.dataset import PretokenizedDataset
+
+class TestPretokenizedDatasetContinuity(unittest.TestCase):
+    def setUp(self):
+        self.test_dir = Path(tempfile.mkdtemp())
+        self.seq_len = 5
+        self.chunk_len = self.seq_len + 1 # 6
+        
+        # Create two files.
+        # File 1: 0, 1, 2, 3 (4 tokens) - too small for one sequence
+        # File 2: 4, 5, 6, 7, 8, 9, 10, 11 (8 tokens)
+        
+        # Combined: 0..11 (12 tokens)
+        # Sequence 1: 0,1,2,3,4,5 (input 0..4, target 1..5)
+        # Sequence 2: 6,7,8,9,10,11 (input 6..10, target 7..11)
+        
+        # Wait, simple streaming uses chunks of size seq_len + 1.
+        # So we need 6 tokens per yield.
+        
+        self.split = "train"
+        
+        f1 = self.test_dir / f"data_{self.split}_001.bin"
+        tokens1 = np.arange(4, dtype=np.uint16)
+        with open(f1, "wb") as f:
+            f.write(tokens1.tobytes())
+            
+        f2 = self.test_dir / f"data_{self.split}_002.bin"
+        tokens2 = np.arange(4, 12, dtype=np.uint16)
+        with open(f2, "wb") as f:
+            f.write(tokens2.tobytes())
+            
+    def tearDown(self):
+        shutil.rmtree(self.test_dir)
+        
+    def test_continuity(self):
+        ds = PretokenizedDataset(self.test_dir, seq_len=self.seq_len, split=self.split)
+        iterator = iter(ds)
+        
+        # First batch should come from merging File 1 and start of File 2
+        # Data: 0,1,2,3 + 4,5
+        x1, y1 = next(iterator)
+        print(f"Batch 1: {x1.tolist()}")
+        self.assertEqual(x1.tolist(), [0, 1, 2, 3, 4])
+        self.assertEqual(y1.tolist(), [1, 2, 3, 4, 5])
+        
+        # Second batch should come from File 2
+        # Remaining in File 2: 6,7,8,9,10,11 (6 tokens left)
+        # This is exactly one chunk.
+        
+        # But wait, my logic was:
+        # 1. needed = 6 - 4 = 2.
+        # 2. Take 2 from File 2 (4,5). Yield.
+        # 3. current_idx becomes 2.
+        # 4. Total len of File 2 is 8. Remaining = 6.
+        # 5. num_chunks = 1.
+        # 6. Yield next chunk (from idx 2 to 8 -> 6,7,8,9,10,11).
+        
+        try:
+            x2, y2 = next(iterator)
+            print(f"Batch 2: {x2.tolist()}")
+            self.assertEqual(x2.tolist(), [6, 7, 8, 9, 10])
+            self.assertEqual(y2.tolist(), [7, 8, 9, 10, 11])
+        except StopIteration:
+            self.fail("Should have yielded a second batch")
+            
+        # Should be empty now
+        with self.assertRaises(StopIteration):
+            next(iterator)
+
+if __name__ == "__main__":
+    unittest.main()