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

astrai/dataset/__init__.py

@@ -9,10 +9,10 @@ from astrai.dataset.storage import (
     H5Store,
     JsonlStore,
     MmapStore,
-    RecordStore,
+    Recordable,
     Store,
     StoreFactory,
-    StreamStore,
+    Streamable,
     detect_format,
 )
 from astrai.serialization import (
@@ -28,8 +28,8 @@ __all__ = [
     "dpo_collate_fn",
     "grpo_collate_fn",
     "Store",
-    "StreamStore",
-    "RecordStore",
+    "Streamable",
+    "Recordable",
     "StoreFactory",
     "H5Store",
     "MmapStore",

astrai/dataset/dataset.py

@@ -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,

astrai/dataset/storage.py

@@ -1,24 +1,37 @@
 """Storage backends for different data formats.
 
-Two access modes are reflected in the class hierarchy:
+Architecture (mixin composition, no diamond inheritance):
 
-- :class:`StreamStore` — ``fetch(begin, end, key)`` slices across
-  concatenated segments.  Used by PT/SFT where data is a long token
-  stream.  ``len(store)`` returns the total token count.
-- :class:`RecordStore` — ``fetch_record(i, key)`` returns the *i*-th
-  record without cross-record concatenation.  Used by DPO/GRPO where
-  each record is an independent training unit.  ``num_records`` returns
+    Store (ABC)              — shared _data/_cum/_offsets bookkeeping
+                               + _normalize() for registering segments
+    Streamable (mixin)       — fetch(begin, end, key) for stream access
+    Recordable (mixin)       — fetch_record(i, key) for record access
+
+    H5Store(Store, Streamable, Recordable)
+    MmapStore(Store, Streamable, Recordable)
+    JsonlStore(Store, Streamable, Recordable)
+
+Each mixin is a stateless trait that relies on ``self._data`` etc.
+provided by :class:`Store`.  Concrete stores mix in whichever access
+modes they support — ``Store`` is the sole base class, so there is no
+diamond inheritance or MRO ambiguity.
+
+Access-mode semantics:
+
+- **Stream** (SEQ/SFT): ``fetch(begin, end, key)`` slices across
+  concatenated segments.  ``len(store)`` returns the total token count.
+- **Record** (DPO/GRPO): ``fetch_record(i, key)`` returns the *i*-th
+  record without cross-record concatenation.  ``num_records`` returns
   the record count.
 
-Both share ``_data`` / ``_cum`` / ``_offsets`` bookkeeping via the
-common :class:`Store` base, which also owns ``_normalize`` for
-registering segments.  Subclasses pick the access mode by inheriting
-from the appropriate base.
+``segments_are_records`` (class attribute on each Store subclass)
+tells ``_normalize`` whether segments are inherently per-record (H5/
+JSONL) or opaque shards (bin).  Record access for bin relies on
+``_offsets`` instead.
 
-:class:`ProcessedStore` composes a :class:`JsonlSource` (raw record
-reader) with a pure ``record -> dict_of_tensors`` processor so that
-DPO/GRPO can tokenise raw JSONL on the fly without a pre-tokenised
-H5/bin file.  This keeps the tokenizer out of the Store base.
+:class:`JsonlStore` supports a lazy mode (``processor=fn``) that keeps
+raw records and defers tokenisation to ``fetch_record`` — used by DPO
+to train directly from a ``.jsonl`` file without a pre-tokenised copy.
 """
 
 import bisect
@@ -112,6 +125,14 @@ class Store(ABC):
     def keys(self) -> List[str]:
         return list(self._data.keys())
 
+    def __len__(self) -> int:
+        """Default: token count (stream semantics).
+
+        Subclasses that are record-only (e.g. lazy JsonlStore) override
+        to return ``self._num_records``.
+        """
+        return self._length
+
     def _normalize(
         self,
         raw: Dict[str, list],
@@ -181,11 +202,13 @@ class Store(ABC):
             self._num_records = 0
 
 
-class StreamStore(Store):
-    """Store exposing stream access: ``fetch(begin, end, key)``."""
+class Streamable:
+    """Mixin: stream access ``fetch(begin, end, key)``.
 
-    def __len__(self) -> int:
-        return self._length
+    No base class — relies on ``self._data``, ``self._cum``,
+    ``self._length`` provided by :class:`Store`.  Used by SEQ/SFT
+    where data is a long token stream.
+    """
 
     def fetch(
         self,
@@ -200,10 +223,10 @@ class StreamStore(Store):
                 f"Index out of bounds: begin={begin}, end={end}, length={self._length}"
             )
         if isinstance(keys, str):
-            return self._fetch_key(keys, begin, end)
-        return {k: self._fetch_key(k, begin, end) for k in keys}
+            return self._fetch_stream_key(keys, begin, end)
+        return {k: self._fetch_stream_key(k, begin, end) for k in keys}
 
-    def _fetch_key(self, key: str, begin: int, end: int) -> Tensor:
+    def _fetch_stream_key(self, key: str, begin: int, end: int) -> Tensor:
         segments = self._data[key]
         cum = self._cum[key]
         seg_start = bisect.bisect_right(cum, begin)
@@ -219,11 +242,12 @@ class StreamStore(Store):
         return results[0] if len(results) == 1 else torch.cat(results, dim=0)
 
 
-class RecordStore(Store):
-    """Mixin exposing record access: ``fetch_record(i, key)``.
+class Recordable:
+    """Mixin: record access ``fetch_record(i, key)``.
 
-    ``__len__`` is **not** overridden — subclasses decide whether
-    ``len()`` returns token count (stream) or record count (record-only).
+    No base class — relies on ``self._data``, ``self._offsets``,
+    ``self._num_records`` provided by :class:`Store`.  Used by
+    DPO/GRPO where each record is an independent training unit.
     """
 
     segments_are_records = True
@@ -237,7 +261,7 @@ class RecordStore(Store):
         index: int,
         keys: Union[str, List[str]],
     ):
-        if not self._data:
+        if not self._data and self._num_records == 0:
             raise RuntimeError("Store not loaded")
         if not 0 <= index < self._num_records:
             raise ValueError(
@@ -265,7 +289,7 @@ class StoreFactory(BaseFactory["Store"]):
 
 
 @StoreFactory.register("h5")
-class H5Store(StreamStore, RecordStore):
+class H5Store(Store, Streamable, Recordable):
     """HDF5-based storage backend (pre-tokenized data).
 
     Each key is stored as a group of per-record datasets (``data_0``,
@@ -281,12 +305,14 @@ class H5Store(StreamStore, RecordStore):
     ``store.num_records`` instead.
     """
 
+    segments_are_records = True
+
     def load(self, path: str, **kwargs):
         self._normalize(load_h5(path))
 
 
 @StoreFactory.register("bin")
-class MmapStore(StreamStore, RecordStore):
+class MmapStore(Store, Streamable, Recordable):
     """Memory-mapped binary storage backend.
 
     Each key is a single .bin file backed by ``np.memmap(mode="r")``.
@@ -301,9 +327,13 @@ class MmapStore(StreamStore, RecordStore):
       ``save_bin(..., record_keys=...)``).  Legacy bin files without
       offsets have ``num_records == 0``.
 
-    ``len(store)`` returns the **token count** (stream semantics).
+    ``segments_are_records`` is ``False`` here (bin segments are
+    contiguous streams, not per-record) — record access is driven
+    purely by ``_offsets``.
     """
 
+    segments_are_records = False
+
     def load(self, path: str, **kwargs):
         self._mmap_refs = []
         root = Path(path)
@@ -375,7 +405,7 @@ class JsonlSource:
 
 
 @StoreFactory.register("jsonl")
-class JsonlStore(StreamStore, RecordStore):
+class JsonlStore(Store, Streamable, Recordable):
     """JSONL reader with two tokenisation modes.
 
     A JSONL dataset is a ``.jsonl`` file or a directory of ``*.jsonl``
@@ -399,6 +429,7 @@ class JsonlStore(StreamStore, RecordStore):
     """
 
     CONFIG_NAME = "dataset_config.json"
+    segments_are_records = True
 
     def __init__(self):
         super().__init__()

tests/data/test_dataset.py

@@ -1,14 +1,16 @@
 import json
 import os
+import tempfile
 
 import numpy as np
 import pytest
 import torch
 
 from astrai.config.preprocess_config import PipelineConfig
-from astrai.dataset.dataset import DatasetFactory, SEQDataset
+from astrai.dataset.dataset import DatasetFactory, SEQDataset, dpo_tokenize
 from astrai.dataset.storage import (
     H5Store,
+    JsonlStore,
     StoreFactory,
     detect_format,
 )
@@ -900,3 +902,166 @@ def test_grpo_multiple_records(base_test_env):
         assert item["rewards"].shape == (G,)
         for g in range(G):
             assert item["responses"][g].shape == item["masks"][g].shape
+
+
+def _write_dpo_jsonl(test_dir, records):
+    """Write a raw DPO JSONL file (no dataset_config.json)."""
+    path = os.path.join(test_dir, "dpo.jsonl")
+    with open(path, "w", encoding="utf-8") as f:
+        for rec in records:
+            f.write(json.dumps(rec, ensure_ascii=False) + "\n")
+    return path
+
+
+def test_dpo_tokenize_pure_function():
+    """dpo_tokenize returns flat lists with correct mask alignment."""
+
+    class FakeTokenizer:
+        def encode(self, text, add_special_tokens=True):
+            return [len(text)] if add_special_tokens else [len(text) + 1]
+
+    record = {"input": "ab", "chosen": "xyz", "rejected": "w"}
+    result = dpo_tokenize(record, FakeTokenizer(), max_len=64, pad_id=0)
+
+    assert set(result.keys()) == {"chosen", "rejected", "chosen_mask", "rejected_mask"}
+    assert len(result["chosen"]) == len(result["chosen_mask"])
+    assert len(result["rejected"]) == len(result["rejected_mask"])
+    assert len(result["chosen"]) == len(result["rejected"])
+
+    assert result["chosen_mask"][0] == 0
+    assert any(m == 1 for m in result["chosen_mask"])
+    assert result["rejected_mask"][0] == 0
+
+
+def test_dpo_tokenize_malformed_record():
+    """dpo_tokenize returns None for missing fields."""
+
+    class FakeTokenizer:
+        def encode(self, text, add_special_tokens=True):
+            return [1]
+
+    assert dpo_tokenize({}, FakeTokenizer()) is None
+    assert dpo_tokenize({"input": "a"}, FakeTokenizer()) is None
+    assert dpo_tokenize({"input": "a", "chosen": "b"}, FakeTokenizer()) is None
+
+
+def test_dpo_jsonl_lazy_load(base_test_env):
+    """DPODataset loads raw JSONL with tokenizer_path → lazy processor."""
+    test_dir = base_test_env["test_dir"]
+    tokenizer_path = _save_test_tokenizer(test_dir, base_test_env["tokenizer"])
+
+    records = [
+        {"input": "Hello", "chosen": "world", "rejected": "earth"},
+        {"input": "Foo", "chosen": "bar", "rejected": "baz"},
+    ]
+    path = _write_dpo_jsonl(test_dir, records)
+
+    ds = DatasetFactory.load(
+        train_type="dpo",
+        load_path=path,
+        window_size=0,
+        tokenizer_path=tokenizer_path,
+    )
+
+    assert len(ds) == 2
+    assert ds.storage.num_records == 2
+    assert ds.storage._processor is not None
+
+    item = ds[0]
+    assert set(item.keys()) == {"chosen", "rejected", "chosen_mask", "rejected_mask"}
+    assert item["chosen"].dtype == torch.long
+    assert item["chosen_mask"].dtype == torch.bool
+    assert item["chosen"].shape == item["chosen_mask"].shape
+    assert item["chosen"].shape == item["rejected"].shape
+
+
+def test_dpo_jsonl_lazy_no_tokenizer():
+    """DPODataset on jsonl without tokenizer_path falls back to eager
+    (which requires dataset_config.json, so it should raise)."""
+    with tempfile.TemporaryDirectory() as d:
+        path = os.path.join(d, "dpo.jsonl")
+        with open(path, "w") as f:
+            f.write(json.dumps({"input": "a", "chosen": "b", "rejected": "c"}) + "\n")
+
+        with pytest.raises(FileNotFoundError, match="dataset_config.json"):
+            DatasetFactory.load(
+                train_type="dpo",
+                load_path=path,
+                window_size=0,
+            )
+
+
+def test_jsonl_store_lazy_len_returns_record_count(base_test_env):
+    """JsonlStore in lazy mode: len() returns record count, not tokens."""
+    test_dir = base_test_env["test_dir"]
+    records = [{"input": str(i), "chosen": "c", "rejected": "r"} for i in range(5)]
+    path = _write_dpo_jsonl(test_dir, records)
+
+    store = JsonlStore()
+    store.load(path, processor=lambda r: {"chosen": torch.tensor([1, 2])})
+
+    assert len(store) == 5
+    assert store.num_records == 5
+
+
+def test_jsonl_store_eager_len_returns_token_count(base_test_env):
+    """JsonlStore in eager mode: num_records reflects per-record count."""
+    test_dir = base_test_env["test_dir"]
+    tokenizer_path = _save_test_tokenizer(test_dir, base_test_env["tokenizer"])
+    data_dir = _write_jsonl_dataset(
+        test_dir,
+        tokenizer_path,
+        [{"text": "hello world"}, {"text": "foo bar"}],
+        config_overrides={
+            "preprocessing": {"max_seq_len": 128, "min_chars": 0},
+            "output": {"position_ids_mode": "none"},
+        },
+    )
+
+    store = JsonlStore()
+    store.load(data_dir)
+
+    assert store.num_records == 2
+    assert len(store.keys) > 0
+
+
+def test_h5_store_dual_mode(base_test_env):
+    """H5Store supports both fetch (stream) and fetch_record (record)."""
+    test_dir = base_test_env["test_dir"]
+
+    seq_length = 64
+    dummy_data = {
+        "chosen": [_rand_seq(seq_length), _rand_seq(seq_length)],
+        "rejected": [_rand_seq(seq_length), _rand_seq(seq_length)],
+    }
+    save_h5(test_dir, "dpo_data", dummy_data)
+
+    store = H5Store()
+    store.load(test_dir)
+
+    assert len(store) == seq_length * 2
+    assert store.num_records == 2
+
+    rec0 = store.fetch_record(0, "chosen")
+    assert rec0.shape == (seq_length,)
+
+    stream = store.fetch(0, 10, "chosen")
+    assert stream.shape == (10,)
+
+
+def test_mmap_store_stream_only_no_offsets(base_test_env):
+    """MmapStore without offsets: num_records == 0, stream works."""
+    test_dir = base_test_env["test_dir"]
+
+    seq_length = 128
+    dummy_data = {"sequence": [_rand_seq(seq_length)]}
+    save_bin(test_dir, dummy_data)
+
+    store = StoreFactory.create("bin")
+    store.load(test_dir)
+
+    assert len(store) == seq_length
+    assert store.num_records == 0
+
+    chunk = store.fetch(0, 32, "sequence")
+    assert chunk.shape == (32,)