refactor: replace diamond inheritance with mixin composition

- StreamStore/RecordStore → Streamable/Recordable (stateless mixins)
- Store is sole base class, no MRO ambiguity
- H5Store/MmapStore/JsonlStore mix in both traits explicitly
- segments_are_records declared per-subclass (H5/Jsonl=True, bin=False)
- Add tests for dpo_tokenize, lazy jsonl, dual-mode H5, stream-only bin
- Remove unused _to_tensor helper
This commit is contained in:
2026-07-18 23:20:41 +08:00
parent b133fc9c07
commit 553a42702d
4 changed files with 251 additions and 47 deletions
+19 -11
View File
@@ -1,4 +1,22 @@
"""Dataset implementations with factory pattern for training."""
"""Dataset implementations with factory pattern for training.
Class hierarchy:
BaseDataset (ABC) — load/validate, owns a Store
├── SEQDataset — stream, next-token prediction (PT)
├── SFTDataset — stream, loss-mask + position_ids
└── RecordDataset — record access, optional processor
├── DPODataset — chosen/rejected pairs
└── GRPODataset — prompt + response group
``RecordDataset`` holds an optional *processor* (pure
``record -> Dict[str, Tensor]`` function). When the backing Store is
a lazy JsonlStore, the processor tokenises on the fly; otherwise it
is ignored and ``fetch_record`` reads pre-tokenised tensors.
``__len__`` returns the sample count (stream: windows, record:
records) so DataLoader and progress bars work uniformly.
"""
from abc import ABC, abstractmethod
from functools import partial
@@ -9,24 +27,14 @@ from torch import Tensor
from torch.utils.data import Dataset
from astrai.dataset.storage import (
RecordStore,
Store,
StoreFactory,
StreamStore,
detect_format,
)
from astrai.factory import BaseFactory
from astrai.tokenize import AutoTokenizer
def _to_tensor(value: list, dtype: Optional[torch.dtype] = None) -> Tensor:
if dtype is not None:
return torch.tensor(value, dtype=dtype)
if value and isinstance(value[0], bool):
return torch.tensor(value, dtype=torch.bool)
return torch.tensor(value, dtype=torch.int32)
def dpo_tokenize(
record: dict,
tokenizer,