feat: add streaming IterableDataset for pretraining

- StreamingSeqDataset yields windows sequentially through each shard
- Shard-level shuffle, distributed and multi-worker shard partitioning
- __len__ returns total window count for scheduler total_steps
- Better OS page-cache locality than random-access map-style datasets
This commit is contained in:
2026-08-08 16:15:10 +08:00
parent 9bff61fb91
commit d7cd69fef5
2 changed files with 124 additions and 0 deletions
+2
View File
@@ -14,6 +14,7 @@ from astrai.dataset.storage import (
Streamable,
detect_format,
)
from astrai.dataset.streaming import StreamingSeqDataset
from astrai.serialization import (
load_bin,
save_bin,
@@ -34,4 +35,5 @@ __all__ = [
"save_bin",
"load_bin",
"RDSampler",
"StreamingSeqDataset",
]