pretrain dataset

scripts/TRAIN.md

@@ -21,6 +21,17 @@ python -B train_tokenizer.py
 python -B prepare_pretrain_dataset.py
 ```
 
+```python
+from litdata import StreamingDataset, StreamingDataLoader, TokensLoader
+
+dataset = StreamingDataset(
+  input_dir='../pretrain-data/',
+  item_loader=TokensLoader(block_size=2048 + 1),
+)
+
+print(len(dataset))
+```
+
 ## Model
 
 ### Pretrain

scripts/pretrain-model.yaml

@@ -77,7 +77,7 @@ train:
 
   # Total number of tokens to train on (type: Optional[int], default: 3000000000000)
   # max_tokens: 3000000000000
-  max_tokens: 8159107755 # 796399 * 2049 * 5
+  max_tokens: 36852166560 # 3597088 * 2049 * 5
 
   # Limits the number of optimizer steps to run. (type: Optional[int], default: null)
   max_steps: