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:
+19
-11
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user