refactor: split Store into StreamStore and RecordStore

- StreamStore: fetch(begin, end, key) for stream access (SEQ/SFT)
- RecordStore: mixin with fetch_record(i, key) for record access
- H5Store/MmapStore/JsonlStore now dual-inherit both (C3 MRO)
- JsonlStore supports lazy mode via processor= (no TokenizeTransform)
- RecordDataset base class holds processor, DPO/GRPO simplified
- dpo_tokenize pure function for on-the-fly JSONL tokenisation
- DatasetFactory builds processor for jsonl+record datasets
- train.py passes tokenizer_path=param_path uniformly
- progress: len(dataset) returns sample count (stream=windows, record=records)
- json no longer auto-detected as jsonl format

astrai/dataset/__init__.py

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

astrai/dataset/dataset.py

@@ -1,18 +1,93 @@
 """Dataset implementations with factory pattern for training."""
 
 from abc import ABC, abstractmethod
-from typing import Dict, List, Optional
+from functools import partial
+from typing import Callable, Dict, List, Optional
 
 import torch
 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,
+    max_len: int = 2048,
+    pad_id: int = 2,
+) -> Optional[dict]:
+    """Tokenize one DPO record into chosen/rejected + masks.
+
+    Pure processor function (HF ``datasets.map`` style):
+    ``record -> dict_of_lists``.  Each value is a flat list of ints/bools.
+
+    No packing, no ``position_ids`` — DPO sequences are independent and
+    the model defaults to ``arange(0, seq_len)``.
+    """
+    inp = record.get("input")
+    chosen_text = record.get("chosen")
+    rejected_text = record.get("rejected")
+    if inp is None or chosen_text is None or rejected_text is None:
+        return None
+
+    in_ids = tokenizer.encode(inp, add_special_tokens=True)
+    ch_ids = tokenizer.encode(chosen_text, add_special_tokens=False)
+    re_ids = tokenizer.encode(rejected_text, add_special_tokens=False)
+
+    full_ch = (in_ids + ch_ids)[:max_len]
+    full_re = (in_ids + re_ids)[:max_len]
+
+    max_record_len = max(len(full_ch), len(full_re))
+    ch_pad = full_ch + [pad_id] * (max_record_len - len(full_ch))
+    re_pad = full_re + [pad_id] * (max_record_len - len(full_re))
+
+    ch_mask = [0] * len(in_ids) + [1] * len(ch_ids)
+    ch_mask = ch_mask[:max_len]
+    ch_mask += [0] * (max_record_len - len(ch_mask))
+    re_mask = [0] * len(in_ids) + [1] * len(re_ids)
+    re_mask = re_mask[:max_len]
+    re_mask += [0] * (max_record_len - len(re_mask))
+
+    return {
+        "chosen": ch_pad,
+        "rejected": re_pad,
+        "chosen_mask": ch_mask,
+        "rejected_mask": re_mask,
+    }
+
+
+def dpo_processor(
+    record: dict,
+    tokenizer,
+    max_len: int = 2048,
+) -> Dict[str, Tensor]:
+    """DPO processor: wraps :func:`dpo_tokenize` and returns tensors."""
+    result = dpo_tokenize(record, tokenizer, max_len=max_len)
+    if result is None:
+        raise ValueError(f"Malformed DPO record: {list(record.keys())}")
+    return {
+        "chosen": torch.tensor(result["chosen"], dtype=torch.int32),
+        "rejected": torch.tensor(result["rejected"], dtype=torch.int32),
+        "chosen_mask": torch.tensor(result["chosen_mask"], dtype=torch.bool),
+        "rejected_mask": torch.tensor(result["rejected_mask"], dtype=torch.bool),
+    }
 
 
 def dpo_collate_fn(batch: List[Dict[str, Tensor]]) -> Dict[str, Tensor]:
@@ -206,6 +281,63 @@ class BaseDataset(Dataset, ABC):
         return (total - 1 - self.window_size) // self.stride + 1
 
 
+class RecordDataset(BaseDataset):
+    """Base class for record-structured datasets (DPO/GRPO).
+
+    Each sample is an independent record — no windowing, stride, or
+    cross-record concatenation.  ``__len__`` returns the record count
+    so progress bars advance per-record.
+
+    A *processor* (pure ``record -> Dict[str, Tensor]`` function) may be
+    supplied for lazy on-the-fly tokenisation of raw JSONL.  The
+    processor is forwarded to ``JsonlStore`` and applied per access;
+    pre-tokenised backends (H5/bin) ignore it.
+    """
+
+    def __init__(
+        self,
+        window_size: int = 0,
+        stride: int = 0,
+        processor: Optional[Callable[[dict], Dict[str, Tensor]]] = None,
+        **kwargs,
+    ):
+        super().__init__(window_size=window_size, stride=stride or window_size)
+        self.processor = processor
+
+    def load(self, load_path: str, storage_type: Optional[str] = None, **kwargs):
+        """Load data from *load_path*.
+
+        Args:
+            load_path: Path to data file or directory.
+            storage_type: Force backend ("h5"/"bin"/"jsonl") or None for
+                auto-detection.
+            **kwargs: Forwarded to ``store.load()``.  When the backend is
+                JSONL and a processor was set, it is passed as
+                ``processor=`` for lazy tokenisation.
+        """
+        if storage_type is None:
+            storage_type = detect_format(load_path)
+        self.storage = StoreFactory.create(storage_type, **kwargs)
+        self._load_path = load_path
+
+        if self.processor is not None:
+            self.storage.load(load_path, processor=self.processor, **kwargs)
+        else:
+            self.storage.load(load_path, **kwargs)
+            self._validate_keys()
+
+    def __len__(self) -> int:
+        if self.storage is None:
+            return 0
+        return self.storage.num_records
+
+    @property
+    def count(self) -> int:
+        if self.storage is None:
+            return 0
+        return self.storage.num_records
+
+
 class DatasetFactory(BaseFactory["BaseDataset"]):
     """Factory class for creating dataset instances.
 
@@ -229,6 +361,8 @@ class DatasetFactory(BaseFactory["BaseDataset"]):
         window_size: int,
         stride: Optional[int] = None,
         storage_type: Optional[str] = None,
+        tokenizer_path: Optional[str] = None,
+        max_len: int = 2048,
         **kwargs,
     ) -> "BaseDataset":
         """Create and load a dataset in one step.
@@ -239,6 +373,11 @@ class DatasetFactory(BaseFactory["BaseDataset"]):
             window_size: Window size for data sampling
             stride: Stride between consecutive samples (default: same as window_size)
             storage_type: Storage type ("h5", "bin", "jsonl") or None for auto-detection
+            tokenizer_path: Path to tokenizer. Used to build an on-the-fly
+                processor when loading raw JSONL with a record dataset
+                (DPO/GRPO).  Ignored for pre-tokenised backends (H5/bin)
+                and for stream datasets (SEQ/SFT).
+            max_len: Max sequence length for the processor.
             **kwargs: Extra arguments forwarded to ``dataset.load()``.
 
         Returns:
@@ -247,11 +386,57 @@ class DatasetFactory(BaseFactory["BaseDataset"]):
         if stride is None:
             stride = window_size
 
-        dataset = cls.create(train_type, window_size, stride)
+        if storage_type is None:
+            storage_type = detect_format(load_path)
+
+        processor = cls._maybe_build_processor(
+            train_type, storage_type, tokenizer_path, max_len
+        )
+
+        dataset = cls.create(train_type, window_size, stride, processor=processor)
         dataset.load(load_path, storage_type=storage_type, **kwargs)
 
         return dataset
 
+    @classmethod
+    def from_store(
+        cls,
+        train_type: str,
+        store: Store,
+        window_size: int = 0,
+        stride: Optional[int] = None,
+    ) -> "BaseDataset":
+        """Create a dataset bound to an already-loaded store.
+
+        The caller is responsible for constructing and loading the store
+        (including any processor).  The dataset simply wraps it.
+        """
+        if stride is None:
+            stride = window_size
+        dataset = cls.create(train_type, window_size, stride)
+        dataset.storage = store
+        return dataset
+
+    @staticmethod
+    def _maybe_build_processor(
+        train_type: str,
+        storage_type: str,
+        tokenizer_path: Optional[str],
+        max_len: int,
+    ) -> Optional[Callable[[dict], Dict[str, Tensor]]]:
+        """Build an on-the-fly tokenisation processor if applicable.
+
+        Only raw JSONL + record datasets (DPO/GRPO) need a processor;
+        pre-tokenised backends (H5/bin) and stream datasets (SEQ/SFT)
+        return ``None`` so no tokenizer is loaded.
+        """
+        if tokenizer_path is None or storage_type != "jsonl":
+            return None
+        if train_type == "dpo":
+            tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
+            return partial(dpo_processor, tokenizer=tokenizer, max_len=max_len)
+        return None
+
 
 @DatasetFactory.register("seq")
 class SEQDataset(BaseDataset):
@@ -301,7 +486,7 @@ class SFTDataset(BaseDataset):
 
 
 @DatasetFactory.register("dpo")
-class DPODataset(BaseDataset):
+class DPODataset(RecordDataset):
     """Record-structured dataset for Direct Preference Optimization.
 
     Each sample is one preference pair (chosen + rejected) and is an
@@ -309,41 +494,21 @@ class DPODataset(BaseDataset):
     concatenation.  This keeps each sequence self-contained so attention
     never leaks across preference pairs.
 
-    Delegates record access to ``Store.fetch_record``, which works with
+    Two loading paths (handled by :class:`RecordDataset`):
-    any storage backend (H5 per-record datasets, bin+offsets memmap, or
-    JSONL on-the-fly tokenization).
-    """
 
-    def __init__(self, window_size: int = 0, stride: int = 0, **kwargs):
+    - **Pre-tokenized** (H5/bin): ``load(path)`` reads per-record tensors,
-        super().__init__(window_size=window_size, stride=stride or window_size)
+      ``__getitem__`` returns them directly.
+    - **Raw JSONL** (``tokenizer_path=...``): builds a lazy processor via
+      :func:`dpo_processor` that tokenises on the fly — no packing, no
+      ``position_ids``.
+    """
 
     @property
     def required_keys(self) -> List[str]:
         return ["chosen", "rejected", "chosen_mask", "rejected_mask"]
 
-    def load(self, load_path: str, storage_type: Optional[str] = None, **kwargs):
+    def make_processor(self, tokenizer, max_len: int):
-        if storage_type is None:
+        return partial(dpo_processor, tokenizer=tokenizer, max_len=max_len)
-            storage_type = detect_format(load_path)
-        self.storage = StoreFactory.create(storage_type, **kwargs)
-        self._load_path = load_path
-        self.storage.load(load_path, **kwargs)
-        self._validate_keys()
-
-    def _validate_keys(self):
-        actual_keys = set(self.storage.keys)
-        missing = [k for k in self.required_keys if k not in actual_keys]
-        if missing:
-            raise KeyError(
-                f"DPODataset requires keys {self.required_keys}, "
-                f"but storage only has {sorted(actual_keys)}. Missing: {missing}"
-            )
-
-    @property
-    def count(self) -> int:
-        return self.storage.num_records
-
-    def __len__(self) -> int:
-        return self.storage.num_records
 
     def __getitem__(self, index: int) -> Dict[str, Tensor]:
         return {
@@ -361,13 +526,11 @@ class DPODataset(BaseDataset):
 
 
 @DatasetFactory.register("grpo")
-class GRPODataset(BaseDataset):
+class GRPODataset(RecordDataset):
     """Dataset for offline Group Relative Policy Optimization.
 
-    Unlike the window-based datasets (SEQ/SFT/DPO), GRPO data is
+    Each sample is one prompt with its group of responses and scalar
-    record-structured: each sample is one prompt with its group of
+    rewards — an independent training unit with no windowing or stride.
-    responses and scalar rewards.  There is no windowing or stride —
-    every record is an independent training unit.
 
     Expected storage layout (produced by JsonlStore or pre-tokenized):
 
@@ -377,37 +540,10 @@ class GRPODataset(BaseDataset):
     - ``rewards``:   List[Tensor]  — one 1-D float tensor (len G) per record
     """
 
-    def __init__(self, window_size: int = 0, stride: int = 0, **kwargs):
-        super().__init__(window_size=window_size, stride=stride or window_size)
-
     @property
     def required_keys(self) -> List[str]:
         return ["prompts", "responses", "masks", "rewards"]
 
-    def load(self, load_path: str, storage_type: Optional[str] = None, **kwargs):
-        if storage_type is None:
-            storage_type = detect_format(load_path)
-        self.storage = StoreFactory.create(storage_type, **kwargs)
-        self._load_path = load_path
-        self.storage.load(load_path, **kwargs)
-        self._validate_keys()
-
-    def _validate_keys(self):
-        actual_keys = set(self.storage.keys)
-        missing = [k for k in self.required_keys if k not in actual_keys]
-        if missing:
-            raise KeyError(
-                f"GRPODataset requires keys {self.required_keys}, "
-                f"but storage only has {sorted(actual_keys)}. Missing: {missing}"
-            )
-
-    @property
-    def count(self) -> int:
-        return self.storage.num_records
-
-    def __len__(self) -> int:
-        return self.storage.num_records
-
     def __getitem__(self, index: int) -> Dict[str, Tensor]:
         prompts = self.storage.fetch_record(index, "prompts")
         responses = self.storage.fetch_record(index, "responses")

astrai/dataset/storage.py

@@ -1,20 +1,24 @@
 """Storage backends for different data formats.
 
-Layers:
+Two access modes are reflected in the class hierarchy:
-  - I/O layer:       save_* / load_* functions, read/write raw files (HDF5/bin)
-                      return Dict[str, List[Tensor]] — format-specific, no state
-  - Store (ABC):     central abstraction, normalizes multi-segment into
-                      Dict[str, List[Tensor]] per key via _normalize(),
-                      fetch() uses bisect across segments — no forced concat
-  - Dataset layer:   BaseDataset owns a Store, only calls store.fetch(begin, end, key)
 
-Key properties:
+- :class:`StreamStore` — ``fetch(begin, end, key)`` slices across
-  - Multi-segment:   segments kept as-is, no forced concatenation — safe for
+  concatenated segments.  Used by PT/SFT where data is a long token
-                      datasets larger than RAM
+  stream.  ``len(store)`` returns the total token count.
-  - Explicit length: _length = min(total elements across keys), set at load,
+- :class:`RecordStore` — ``fetch_record(i, key)`` returns the *i*-th
-                      __len__ returns O(1)
+  record without cross-record concatenation.  Used by DPO/GRPO where
-  - Zero-copy mmap:  MmapStore wraps np.memmap(mode="r"), all DataLoader
+  each record is an independent training unit.  ``num_records`` returns
-                       workers share OS page-cache pages
+  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.
+
+: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.
 """
 
 import bisect
@@ -23,7 +27,7 @@ import json
 import logging
 from abc import ABC, abstractmethod
 from pathlib import Path
-from typing import Dict, List, Optional, Union
+from typing import Callable, Dict, List, Optional, Union
 
 import torch
 from torch import Tensor
@@ -46,7 +50,7 @@ def detect_format(load_path: str) -> str:
         load_path: Directory or file path
 
     Returns:
-        Format string ("h5", "bin", or "jsonl")
+        Format string ("h5", "bin", "jsonl", or "processed")
 
     Raises:
         FileNotFoundError: If no supported data files are found
@@ -79,31 +83,16 @@ def detect_format(load_path: str) -> str:
     ]
     if jsonl_files:
         return "jsonl"
-    json_files = [
-        Path(p) for p in glob.glob(str(root / "**" / "*.json"), recursive=True)
-    ]
-    if json_files:
-        return "jsonl"
     raise FileNotFoundError(f"No supported data files found at {load_path}")
 
 
 class Store(ABC):
-    """String keys -> segmented tensors with two access modes.
+    """Common base for all storage backends.
 
-    Stream mode (SEQ/SFT):
+    Owns the shared ``_data`` / ``_cum`` / ``_offsets`` bookkeeping and
-        ``fetch(begin, end, keys)`` slices across concatenated segments,
+    the ``_normalize`` entry point used by tensor-backed subclasses.
-        transparently ``torch.cat``-ing across segment boundaries.
+    Does **not** expose an access API — that is the job of
-        ``len(store)`` returns total token count.
+    :class:`StreamStore` and :class:`RecordStore`.
-
-    Record mode (DPO/GRPO):
-        ``fetch_record(index, keys)`` returns the i-th record without
-        cross-record concatenation.  ``num_records`` returns the record
-        count.  Backed by either per-record segment lists (H5/JSONL) or
-        a single concatenated segment plus per-record offsets (bin).
-
-    Subclasses declare ``segments_are_records`` to indicate whether their
-    segments are inherently per-record (H5/JSONL) or opaque shards (bin).
-    This is a format-level property, not a per-call decision.
     """
 
     segments_are_records: bool = False
@@ -116,94 +105,13 @@ class Store(ABC):
         self._num_records: int = 0
 
     @abstractmethod
-    def load(self, path: str) -> None:
+    def load(self, path: str, **kwargs) -> None:
         raise NotImplementedError
 
     @property
     def keys(self) -> List[str]:
         return list(self._data.keys())
 
-    def __len__(self) -> int:
-        return self._length
-
-    def fetch(
-        self,
-        begin: int,
-        end: int,
-        keys: Union[str, List[str]],
-    ):
-        if not self._data:
-            raise RuntimeError("Store not loaded")
-        if not (0 <= begin < self._length and 0 <= end <= self._length):
-            raise ValueError(
-                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}
-
-    def _fetch_key(self, key: str, begin: int, end: int) -> Tensor:
-        """Fetch slice [begin, end) across potentially multiple segments."""
-        segments = self._data[key]
-        cum = self._cum[key]
-        seg_start = bisect.bisect_right(cum, begin)
-        seg_end = bisect.bisect_left(cum, end)
-
-        results = []
-        for i in range(seg_start, seg_end + 1):
-            prev = cum[i - 1] if i > 0 else 0
-            s = max(begin - prev, 0)
-            e = min(end - prev, segments[i].shape[0])
-            results.append(segments[i][s:e])
-
-        return results[0] if len(results) == 1 else torch.cat(results, dim=0)
-
-    @property
-    def num_records(self) -> int:
-        return self._num_records
-
-    def fetch_record(
-        self,
-        index: int,
-        keys: Union[str, List[str]],
-    ):
-        """Fetch the *index*-th record without cross-record concatenation.
-
-        Returns a tensor (flat key) or ``List[Tensor]`` (nested key such as
-        GRPO ``responses``).
-        """
-        if not self._data:
-            raise RuntimeError("Store not loaded")
-        if not 0 <= index < self._num_records:
-            raise ValueError(
-                f"Record index out of bounds: {index}, num_records={self._num_records}"
-            )
-        if isinstance(keys, str):
-            return self._fetch_record_key(keys, index)
-        return {k: self._fetch_record_key(k, index) for k in keys}
-
-    def _fetch_record_key(self, key: str, index: int):
-        """Return the *index*-th record for *key*.
-
-        Two storage layouts are supported:
-
-        - **bin + offsets**: ``_data[key]`` is ``[single_long_segment]``;
-          ``_offsets[key]`` holds cumulative per-record offsets.  The record
-          is sliced as ``segment[offsets[i]:offsets[i+1]]``.
-        - **h5 / jsonl**: ``_data[key]`` is ``[t0, t1, ...]`` with one tensor
-          (or nested list of tensors for GRPO) per record.  Direct indexing.
-        """
-        offsets = self._offsets.get(key)
-        if offsets:
-            start = offsets[index]
-            end = (
-                offsets[index + 1]
-                if index + 1 < len(offsets)
-                else self._data[key][0].shape[0]
-            )
-            return self._data[key][0][start:end]
-        return self._data[key][index]
-
     def _normalize(
         self,
         raw: Dict[str, list],
@@ -212,17 +120,18 @@ class Store(ABC):
         """Register segments and pre-compute indices for both access modes.
 
         Stream mode: ``_cum[key]`` accumulates per-segment lengths so
-        ``_fetch_key`` can bisect across segments without concatenation.
+        ``StreamStore._fetch_key`` can bisect across segments without
+        concatenation.
 
-        Record mode: if *offsets* is provided (bin layout), ``_offsets[key]``
+        Record mode: if *offsets* is provided (bin layout),
-        stores cumulative per-record offsets into the single concatenated
+        ``_offsets[key]`` stores cumulative per-record offsets into the
-        segment.  Otherwise, when ``segments_are_records`` is True
+        single concatenated segment.  Otherwise, when
-        (H5/JSONL), ``_data[key]`` is a per-record list and
+        ``segments_are_records`` is True (H5/JSONL), ``_data[key]`` is a
-        ``fetch_record`` indexes it directly.
+        per-record list and ``fetch_record`` indexes it directly.
 
-        Nested keys (GRPO ``responses``/``masks`` as ``List[List[Tensor]]``)
+        Nested keys (GRPO ``responses``/``masks`` as
-        are stored as-is and excluded from both cumulative bookkeepings —
+        ``List[List[Tensor]]``) are stored as-is and excluded from both
-        they are only accessed record-by-record.
+        cumulative bookkeepings — they are only accessed record-by-record.
         """
         flat_lengths = []
         for key, tensors in raw.items():
@@ -231,7 +140,6 @@ class Store(ABC):
                 self._cum[key] = []
                 flat_lengths.append(0)
                 continue
-            # Skip nested lists (GRPO responses/masks) — record-level access
             if isinstance(tensors[0], list):
                 self._cum[key] = []
                 continue
@@ -244,9 +152,6 @@ class Store(ABC):
             flat_lengths.append(cum[-1] if cum else 0)
         self._length = min(flat_lengths) if flat_lengths else 0
 
-        # Record-mode offsets (bin layout).  Only valid when each key is a
-        # single concatenated segment — multi-shard bin + offsets is not
-        # supported (merge shards or use H5/JSONL instead).
         valid_offsets: Dict[str, List[int]] = {}
         if offsets:
             for key, off in offsets.items():
@@ -276,53 +181,130 @@ class Store(ABC):
             self._num_records = 0
 
 
-class StoreFactory(BaseFactory["Store"]):
+class StreamStore(Store):
-    """Factory for creating Store instances by type name.
+    """Store exposing stream access: ``fetch(begin, end, key)``."""
 
-    Example::
+    def __len__(self) -> int:
+        return self._length
 
-        @StoreFactory.register("custom")
+    def fetch(
-        class CustomStore(Store):
+        self,
-            ...
+        begin: int,
-    """
+        end: int,
+        keys: Union[str, List[str]],
+    ):
+        if not self._data:
+            raise RuntimeError("Store not loaded")
+        if not (0 <= begin < self._length and 0 <= end <= self._length):
+            raise ValueError(
+                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}
+
+    def _fetch_key(self, key: str, begin: int, end: int) -> Tensor:
+        segments = self._data[key]
+        cum = self._cum[key]
+        seg_start = bisect.bisect_right(cum, begin)
+        seg_end = bisect.bisect_left(cum, end)
+
+        results = []
+        for i in range(seg_start, seg_end + 1):
+            prev = cum[i - 1] if i > 0 else 0
+            s = max(begin - prev, 0)
+            e = min(end - prev, segments[i].shape[0])
+            results.append(segments[i][s:e])
+
+        return results[0] if len(results) == 1 else torch.cat(results, dim=0)
 
 
-@StoreFactory.register("h5")
+class RecordStore(Store):
-class H5Store(Store):
+    """Mixin exposing record access: ``fetch_record(i, key)``.
-    """HDF5-based storage backend (pre-tokenized data).
 
-    Each key is stored as a group of per-record datasets (``data_0``,
+    ``__len__`` is **not** overridden — subclasses decide whether
-    ``data_1``, …), so record mode indexes ``_data[key]`` directly.
+    ``len()`` returns token count (stream) or record count (record-only).
-    Stream mode concatenates across records via ``_cum``.
     """
 
     segments_are_records = True
 
-    def load(self, path: str):
+    @property
+    def num_records(self) -> int:
+        return self._num_records
+
+    def fetch_record(
+        self,
+        index: int,
+        keys: Union[str, List[str]],
+    ):
+        if not self._data:
+            raise RuntimeError("Store not loaded")
+        if not 0 <= index < self._num_records:
+            raise ValueError(
+                f"Record index out of bounds: {index}, num_records={self._num_records}"
+            )
+        if isinstance(keys, str):
+            return self._fetch_record_key(keys, index)
+        return {k: self._fetch_record_key(k, index) for k in keys}
+
+    def _fetch_record_key(self, key: str, index: int):
+        offsets = self._offsets.get(key)
+        if offsets:
+            start = offsets[index]
+            end = (
+                offsets[index + 1]
+                if index + 1 < len(offsets)
+                else self._data[key][0].shape[0]
+            )
+            return self._data[key][0][start:end]
+        return self._data[key][index]
+
+
+class StoreFactory(BaseFactory["Store"]):
+    """Factory for creating Store instances by type name."""
+
+
+@StoreFactory.register("h5")
+class H5Store(StreamStore, RecordStore):
+    """HDF5-based storage backend (pre-tokenized data).
+
+    Each key is stored as a group of per-record datasets (``data_0``,
+    ``data_1``, …).  Supports both access modes:
+
+    - **Stream** (``fetch(begin, end, key)``): concatenates across
+      records via ``_cum`` — used by SEQ/SFT where data is a token stream.
+    - **Record** (``fetch_record(i, key)``): indexes ``_data[key]``
+      directly — used by DPO/GRPO where each record is independent.
+
+    ``len(store)`` returns the **token count** (stream semantics) so
+    SEQ/SFT windowing works.  Record-only code uses
+    ``store.num_records`` instead.
+    """
+
+    def load(self, path: str, **kwargs):
         self._normalize(load_h5(path))
 
 
 @StoreFactory.register("bin")
-class MmapStore(Store):
+class MmapStore(StreamStore, RecordStore):
     """Memory-mapped binary storage backend.
 
     Each key is a single .bin file backed by ``np.memmap(mode="r")``.
     No per-process memory duplication — all DataLoader workers share the
     same OS page-cache pages.
 
-    When ``meta.json`` contains per-record ``offsets`` for a key (written
+    Supports both access modes:
-    via ``save_bin(..., record_keys=...)``), record-mode access slices
-    individual records from the concatenated memmap.  Legacy bin files
-    without offsets only support stream mode.
 
-    Format on disk::
+    - **Stream** (``fetch(begin, end, key)``): always available.
+    - **Record** (``fetch_record(i, key)``): only when ``meta.json``
+      contains per-record ``offsets`` (written via
+      ``save_bin(..., record_keys=...)``).  Legacy bin files without
+      offsets have ``num_records == 0``.
 
-        data_root/
+    ``len(store)`` returns the **token count** (stream semantics).
-          meta.json          # {key: {shape, dtype, offsets?}, ...}
-          <key>.bin          # raw numpy array, one per key
     """
 
-    def load(self, path: str):
+    def load(self, path: str, **kwargs):
         self._mmap_refs = []
         root = Path(path)
         all_raw: Dict[str, List[Tensor]] = {}
@@ -348,50 +330,32 @@ class MmapStore(Store):
             self._mmap_refs.extend(tensors)
 
 
-@StoreFactory.register("jsonl")
+class JsonlSource:
-class JsonlStore(Store):
+    """Read raw JSON records from a ``.jsonl`` file or directory.
-    """JSONL reader with pluggable tokenization transform.
 
-    A JSONL dataset directory contains ``*.jsonl`` files plus a
+    A thin reader used by :class:`JsonlStore` in processor mode — holds
-    ``dataset_config.json`` describing the tokenization pipeline.
+    no tokenizer, performs no tokenisation, just yields dicts.
-
-    Responsibilities are split across two layers:
-
-    - **Reader** (this class): reads raw JSON records from disk.
-    - **Transform** (:class:`~astrai.preprocessing.transform.TokenizeTransform`):
-      tokenizes records into per-key tensors.  When not supplied explicitly,
-      a default transform is built from ``dataset_config.json`` so that
-      existing callers keep working without changes.
-
-    The Store itself never imports the tokenizer — the dependency lives
-    in the Transform layer.
     """
 
-    CONFIG_NAME = "dataset_config.json"
+    def __init__(self, path: str):
-    segments_are_records = True
+        self.path = Path(path)
+        self._records: Optional[List[dict]] = None
 
-    def load(self, path: str, transform=None, **kwargs):
+    def load(self) -> List[dict]:
-        root = Path(path)
+        if self._records is None:
-        records = self._read_records(root)
+            self._records = self._read(self.path)
-
+        return self._records
-        if transform is None:
-            config_path = root / self.CONFIG_NAME
-            if not config_path.exists():
-                raise FileNotFoundError(
-                    f"JSONL dataset config not found: {config_path}. "
-                    f"Expected {self.CONFIG_NAME} alongside *.jsonl files, "
-                    f"or pass an explicit transform."
-                )
-            transform = TokenizeTransform.from_config_file(str(config_path))
-
-        raw = transform.apply(records)
-        self._normalize(raw)
 
     @staticmethod
-    def _read_records(root: Path) -> List[dict]:
+    def _read(root: Path) -> List[dict]:
+        if root.is_file():
+            return JsonlSource._read_file(root)
+        return JsonlSource._read_dir(root)
+
+    @staticmethod
+    def _read_file(path: Path) -> List[dict]:
         records: List[dict] = []
-        for jsonl_path in sorted(root.glob("*.jsonl")):
+        with open(path, "r", encoding="utf-8") as f:
-            with open(jsonl_path, "r", encoding="utf-8") as f:
             for line in f:
                 line = line.strip()
                 if not line:
@@ -399,20 +363,97 @@ class JsonlStore(Store):
                 try:
                     records.append(json.loads(line))
                 except json.JSONDecodeError:
-                        logger.warning(
+                    logger.warning("Failed to parse JSON line in %s, skipping", path)
-                            "Failed to parse JSON line in %s, skipping", jsonl_path
-                        )
-        for json_path in sorted(root.glob("*.json")):
-            if json_path.name == JsonlStore.CONFIG_NAME:
-                continue
-            with open(json_path, "r", encoding="utf-8") as f:
-                try:
-                    data = json.load(f)
-                except json.JSONDecodeError:
-                    logger.warning("Failed to parse JSON file %s, skipping", json_path)
-                    continue
-            if isinstance(data, list):
-                records.extend(data)
-            elif isinstance(data, dict):
-                records.append(data)
         return records
+
+    @staticmethod
+    def _read_dir(root: Path) -> List[dict]:
+        records: List[dict] = []
+        for jsonl_path in sorted(root.glob("*.jsonl")):
+            records.extend(JsonlSource._read_file(jsonl_path))
+        return records
+
+
+@StoreFactory.register("jsonl")
+class JsonlStore(StreamStore, RecordStore):
+    """JSONL reader with two tokenisation modes.
+
+    A JSONL dataset is a ``.jsonl`` file or a directory of ``*.jsonl``
+    files plus (optionally) a ``dataset_config.json`` describing the
+    tokenization pipeline.
+
+    Two modes, selected at :meth:`load` time:
+
+    - **Eager** (default): applies a :class:`TokenizeTransform` to every
+      record at load time and registers per-key tensors via
+      ``_normalize``.  Both stream (``fetch``) and record
+      (``fetch_record``) access work — stream concatenates across
+      records via ``_cum``, record indexes directly.
+    - **Lazy** (``processor=fn`` given): keeps raw records and defers
+      tokenisation to ``fetch_record``.  Only record access works.
+      Used by DPO/GRPO where each record is independent.
+
+    ``len(store)`` returns the **token count** (stream semantics) in
+    eager mode so SEQ/SFT windowing works; returns the **record count**
+    in lazy mode where stream access is unavailable.
+    """
+
+    CONFIG_NAME = "dataset_config.json"
+
+    def __init__(self):
+        super().__init__()
+        self._source: Optional[JsonlSource] = None
+        self._processor: Optional[Callable[[dict], Dict[str, Tensor]]] = None
+        self._keys_cache: Optional[List[str]] = None
+
+    def __len__(self) -> int:
+        if self._processor is not None:
+            return self._num_records
+        return self._length
+
+    def load(self, path: str, transform=None, processor=None, **kwargs):
+        self._source = JsonlSource(path)
+        records = self._source.load()
+
+        if processor is not None:
+            self._processor = processor
+            self._num_records = len(records)
+            return
+
+        if transform is None:
+            root = Path(path)
+            config_path = root / self.CONFIG_NAME if root.is_dir() else None
+            if config_path is None or not config_path.exists():
+                raise FileNotFoundError(
+                    f"JSONL dataset config not found. Expected "
+                    f"{self.CONFIG_NAME} alongside *.jsonl files, pass an "
+                    f"explicit transform, or pass processor= for lazy "
+                    f"on-the-fly tokenisation."
+                )
+            transform = TokenizeTransform.from_config_file(str(config_path))
+
+        transformed = transform.apply(records)
+        self._normalize(transformed)
+
+    @property
+    def keys(self) -> List[str]:
+        if self._processor is not None:
+            if self._keys_cache is None and self._num_records > 0:
+                sample = self._processor(self._source.load()[0])
+                self._keys_cache = list(sample.keys())
+            return self._keys_cache or []
+        return list(self._data.keys())
+
+    def fetch_record(self, index: int, keys: Union[str, List[str]]):
+        if self._processor is not None:
+            if not 0 <= index < self._num_records:
+                raise ValueError(
+                    f"Record index out of bounds: {index}, "
+                    f"num_records={self._num_records}"
+                )
+            record = self._source.load()[index]
+            data = self._processor(record)
+            if isinstance(keys, str):
+                return data[keys]
+            return {k: data[k] for k in keys}
+        return super().fetch_record(index, keys)

scripts/tools/train.py

@@ -460,6 +460,7 @@ def train(
         load_path=data_root_path,
         window_size=window_size,
         stride=stride,
+        tokenizer_path=param_path,
     )
 
     optimizer_fn = partial(

tests/data/test_dataset.py

@@ -460,12 +460,13 @@ def test_dataset_load_explicit_storage_type(base_test_env):
 
 
 def _write_json_dataset(test_dir, tokenizer_path, records, config_overrides=None):
-    """Write JSON (not JSONL) dataset — array of objects."""
+    """Write JSONL dataset — one JSON object per line."""
     data_dir = os.path.join(test_dir, "json_data")
     os.makedirs(data_dir, exist_ok=True)
 
-    with open(os.path.join(data_dir, "data.json"), "w", encoding="utf-8") as f:
+    with open(os.path.join(data_dir, "data.jsonl"), "w", encoding="utf-8") as f:
-        json.dump(records, f, ensure_ascii=False)
+        for rec in records:
+            f.write(json.dumps(rec, ensure_ascii=False) + "\n")
 
     config = {
         "tokenizer_path": tokenizer_path,
@@ -544,7 +545,7 @@ def test_json_store_no_tokenizer_path(base_test_env):
     # Save tokenizer files directly in the dataset directory
     tokenizer.save_pretrained(data_dir)
 
-    # Write .json data
+    # Write .jsonl data
     records = [
         {
             "messages": [
@@ -553,8 +554,9 @@ def test_json_store_no_tokenizer_path(base_test_env):
             ]
         }
     ]
-    with open(os.path.join(data_dir, "data.json"), "w", encoding="utf-8") as f:
+    with open(os.path.join(data_dir, "data.jsonl"), "w", encoding="utf-8") as f:
-        json.dump(records, f, ensure_ascii=False)
+        for rec in records:
+            f.write(json.dumps(rec, ensure_ascii=False) + "\n")
 
     # dataset_config.json WITHOUT tokenizer_path
     config = {