ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor: move sample-id indexing from dataset to store

Browse Source 
- Store owns window_size/stride and __getitem__/__len__/sample_window
- Dataset classes become thin delegators binding a Store to a train-type key mapping
- Drop BaseDataset.get_index and the RecordDataset中间类 (window死代码)
- DatasetFactory forces window_size=0 for record datasets so record semantics never get window-tainted
- token_count/num_records split the legacy len() semantics (raw stream length vs record count)
- Update tests to the new .store/.token_count API and window/record mode switching
This commit is contained in:
ViperEkura
2026-07-19 16:02:50 +08:00
parent 7d478a54db
commit 663ef900fc
3 changed files with 468 additions and 401 deletions
Download Patch File Download Diff File Expand all files Collapse all files
astrai/dataset/storage.py
+257 -111
View File
@@ -1,11 +1,14 @@
"""Storage backends for different data formats.
Architecture (mixin composition, no diamond inheritance):
Architecture (composition over inheritance):
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
Store (ABC) — owns _data/_cum/_offsets bookkeeping
+ window_size/stride for sample-id
indexing. __getitem__/__len__ produce
the smallest iterable unit so Dataset
classes are pure delegators.
Streamable (mixin) — raw token slice fetch(begin, end, keys)
Recordable (mixin) — raw record slice fetch_record(idx, keys)
H5Store(Store, Streamable, Recordable)
MmapStore(Store, Streamable, Recordable)
@@ -13,16 +16,24 @@ Architecture (mixin composition, no diamond inheritance):
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.
primitives they support — ``Store`` is the sole base class, so there is
no diamond inheritance or MRO ambiguity.
Access-mode semantics:
Sample-id indexing lives on :class:`Store`, not on the dataset:
- **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.
- **Stream mode** (``window_size > 0``): ``len(store)`` returns the number
of ``(window_size, stride)`` windows that fit in the token river;
``store[i]`` returns the *i*-th window as a dict of per-key tensors;
``store.sample_window(i)`` exposes the underlying ``(begin, end)``
token slice for callers (e.g. next-token trainers) that need a +1
shifted companion window.
- **Record mode** (``num_records > 0``): ``len(store)`` returns the
record count; ``store[i]`` returns the *i*-th record dict.
Raw token/record access via :meth:`fetch` / :meth:`fetch_record`
remains available for low-level callers that want explicit index
control. ``store.token_count`` is the total stream token count (what
``len(store)`` used to mean in the legacy stream-only API).
``segments_are_records`` (class attribute on each Store subclass)
tells ``_normalize`` whether segments are inherently per-record (H5/
@@ -40,7 +51,7 @@ import json
import logging
from abc import ABC, abstractmethod
from pathlib import Path
from typing import Callable, Dict, List, Optional, Union
from typing import Callable, Dict, List, Optional, Tuple, Union
import torch
from torch import Tensor
@@ -102,20 +113,46 @@ def detect_format(load_path: str) -> str:
class Store(ABC):
"""Common base for all storage backends.
Owns the shared ``_data`` / ``_cum`` / ``_offsets`` bookkeeping and
the ``_normalize`` entry point used by tensor-backed subclasses.
Does **not** expose an access API — that is the job of
:class:`StreamStore` and :class:`RecordStore`.
A Store owns both its data layout AND its sample-id → token/record
index translation. Datasets are thin wrappers that bind a Store
to a particular train-type's key mapping; they never know about
window/stride math.
Two iteration modes:
- **Stream** (``window_size > 0``): data is treated as one long
token river. ``len(store)`` returns the number of windows;
``store[i]`` slices every stream-compatible key to window ``i``;
``store.sample_window(i)`` returns the ``(begin, end)`` token
slice for callers needing a +1 shifted companion window.
- **Record** (``num_records > 0``): data is per-record.
``len(store)`` returns ``num_records``; ``store[i]`` returns
the *i*-th record as a dict.
Raw token slicing is still available via :meth:`fetch` (mixed in
by :class:`Streamable`) when a store has stream support configured.
Raw record slicing via :meth:`fetch_record` (mixed in by
:class:`Recordable`) when a store has record support.
``token_count`` exposes the raw total stream length — this is what
``len(store)`` returned in the legacy stream-only API and what
stream-bound ``fetch`` uses for its bounds check.
"""
segments_are_records: bool = False
def __init__(self):
def __init__(
self,
window_size: int = 0,
stride: Optional[int] = None,
):
self._data: Dict[str, List[Tensor]] = {}
self._cum: Dict[str, List[int]] = {}
self._offsets: Dict[str, List[int]] = {}
self._length: int = 0
self._num_records: int = 0
self._window_size: int = int(window_size)
self._stride: int = int(stride) if stride is not None else int(window_size)
@abstractmethod
def load(self, path: str, **kwargs) -> None:
@@ -125,14 +162,98 @@ class Store(ABC):
def keys(self) -> List[str]:
return list(self._data.keys())
def __len__(self) -> int:
"""Default: token count (stream semantics).
@property
def window_size(self) -> int:
return self._window_size
Subclasses that are record-only (e.g. lazy JsonlStore) override
to return ``self._num_records``.
@property
def stride(self) -> int:
return self._stride
@property
def token_count(self) -> int:
"""Total tokens across all stream segments.
Useful for the bounds-checked raw :meth:`fetch` and as the
legacy ``len(store)`` value.
"""
return self._length
@property
def num_records(self) -> int:
"""Number of records available via :meth:`fetch_record`.
Non-zero only when the backing layout provides per-record
indexing (H5/JSONL segments or bin ``_offsets``).
"""
return self._num_records
@property
def num_samples(self) -> int:
"""Number of items produced by ``__getitem__``.
Stream-mode wins when ``window_size > 0`` and there are tokens
to slice; otherwise falls back to ``num_records``.
"""
if self._window_size > 0 and self._length > 0:
total = self._length
w = self._window_size
if total <= w:
return 0
return (total - 1 - w) // self._stride + 1
return self._num_records
def __len__(self) -> int:
return self.num_samples
def __getitem__(self, index: int) -> Dict[str, Tensor]:
if index < 0:
index += self.num_samples
if not 0 <= index < self.num_samples:
raise IndexError(
f"Store index out of range: {index}, num_samples={self.num_samples}"
)
if self._window_size > 0 and self._length > 0:
begin, end = self.sample_window(index)
keys = self._stream_keys()
return {k: self.fetch(begin, end, k) for k in keys}
return self.fetch_record(index, self._record_keys())
def sample_window(self, index: int) -> Tuple[int, int]:
"""Return ``(begin, end)`` token positions for stream sample *index*.
The clipped tail keeps the last reachable window inside the
token river instead of overshooting. Caller is responsible
for staying within :attr:`num_samples`: an out-of-range index
raises ``IndexError``.
"""
if self._window_size <= 0:
raise RuntimeError("sample_window() requires window_size > 0 (stream mode)")
if self._window_size <= 0 or self._length <= self._window_size:
raise IndexError(
f"Data too short for window: token_count={self._length}, "
f"window_size={self._window_size}"
)
if not 0 <= index < self.num_samples:
raise IndexError(
f"Sample index out of range: {index}, num_samples={self.num_samples}"
)
total = self._length
begin = min(index * self._stride, total - 1 - self._window_size)
end = min(begin + self._window_size, total - 1)
return begin, end
def _stream_keys(self) -> List[str]:
out: List[str] = []
for k, tensors in self._data.items():
if tensors and isinstance(tensors[0], list):
continue
out.append(k)
return out
def _record_keys(self) -> List[str]:
return list(self._data.keys())
def _normalize(
self,
raw: Dict[str, list],
@@ -141,14 +262,14 @@ class Store(ABC):
"""Register segments and pre-compute indices for both access modes.
Stream mode: ``_cum[key]`` accumulates per-segment lengths so
``StreamStore._fetch_key`` can bisect across segments without
concatenation.
``Streamable._fetch_stream_key`` can bisect across segments
without concatenation.
Record mode: if *offsets* is provided (bin layout),
``_offsets[key]`` stores cumulative per-record offsets into the
single concatenated segment. Otherwise, when
``segments_are_records`` is True (H5/JSONL), ``_data[key]`` is a
per-record list and ``fetch_record`` indexes it directly.
``segments_are_records`` is True (H5/JSONL), ``_data[key]`` is
a per-record list and ``fetch_record`` indexes it directly.
Nested keys (GRPO ``responses``/``masks`` as
``List[List[Tensor]]``) are stored as-is and excluded from both
@@ -194,7 +315,7 @@ class Store(ABC):
elif self.segments_are_records:
per_record_counts = []
for key, tensors in self._data.items():
if not tensors or isinstance(tensors[0], list):
if tensors and isinstance(tensors[0], list):
continue
per_record_counts.append(len(tensors))
self._num_records = min(per_record_counts) if per_record_counts else 0
@@ -203,11 +324,13 @@ class Store(ABC):
class Streamable:
"""Mixin: stream access ``fetch(begin, end, key)``.
"""Mixin granting raw token-stream access via :meth:`fetch`.
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.
Stateless trait relying on ``self._data``, ``self._cum``,
``self._length`` maintained by :class:`Store`. Stream mode is
active when the owning store has ``window_size > 0``; for stores
that can also serve record access (H5/JSONL/bin+offsets), the
``fetch_record`` API from :class:`Recordable` is used instead.
"""
def fetch(
@@ -216,72 +339,75 @@ class Streamable:
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_stream_key(keys, begin, end)
return {k: self._fetch_stream_key(k, begin, end) for k in keys}
return _stream_fetch(self, begin, end, keys)
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)
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])
def _stream_fetch(self, begin: int, end: int, keys: Union[str, List[str]]):
if not getattr(self, "_data", None):
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 _fetch_stream_key(self, keys, begin, end)
return {k: _fetch_stream_key(self, k, begin, end) for k in keys}
return results[0] if len(results) == 1 else torch.cat(results, dim=0)
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)
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)
class Recordable:
"""Mixin: record access ``fetch_record(i, key)``.
"""Mixin granting raw record access via :meth:`fetch_record`.
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.
Stateless trait relying on ``self._data``, ``self._offsets``,
``self._num_records`` maintained by :class:`Store`.
"""
segments_are_records = True
@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 and self._num_records == 0:
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}
return _record_fetch(self, index, 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]
def _record_fetch(self, index: int, keys: Union[str, List[str]]):
if not getattr(self, "_data", None) and self._num_records == 0:
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 _fetch_record_key(self, keys, index)
return {k: _fetch_record_key(self, 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"]):
@@ -295,18 +421,22 @@ class H5Store(Store, Streamable, Recordable):
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.
- **Stream**: ``fetch(begin, end, key)`` and ``store[i]`` slice
across concatenated records via ``_cum`` — used by SEQ/SFT.
- **Record**: ``fetch_record(i, key)`` and ``store[i]`` (when
``window_size == 0``) index ``_data[key]`` directly — used by
DPO/GRPO.
"""
segments_are_records = True
def __init__(
self,
window_size: int = 0,
stride: Optional[int] = None,
):
super().__init__(window_size=window_size, stride=stride)
def load(self, path: str, **kwargs):
self._normalize(load_h5(path))
@@ -321,11 +451,12 @@ class MmapStore(Store, Streamable, Recordable):
Supports both access modes:
- **Stream** (``fetch(begin, end, key)``): always available.
- **Stream**: always available via :meth:`fetch`.
- **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``.
offsets have ``num_records == 0`` and ``len(store)`` reflects the
windowed sample count when ``window_size > 0``.
``segments_are_records`` is ``False`` here (bin segments are
contiguous streams, not per-record) — record access is driven
@@ -334,6 +465,14 @@ class MmapStore(Store, Streamable, Recordable):
segments_are_records = False
def __init__(
self,
window_size: int = 0,
stride: Optional[int] = None,
):
super().__init__(window_size=window_size, stride=stride)
self._mmap_refs: List[Tensor] = []
def load(self, path: str, **kwargs):
self._mmap_refs = []
root = Path(path)
@@ -416,32 +555,26 @@ class JsonlStore(Store, Streamable, Recordable):
- **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.
``_normalize``. Both ``fetch`` (stream) and ``fetch_record``
(record) work.
- **Lazy** (``processor=fn`` passed): keeps raw records and defers
tokenisation to ``fetch_record``. Only record access works —
``len(store)`` returns ``num_records``; stream primitives raise.
"""
CONFIG_NAME = "dataset_config.json"
segments_are_records = True
def __init__(self):
super().__init__()
def __init__(
self,
window_size: int = 0,
stride: Optional[int] = None,
):
super().__init__(window_size=window_size, stride=stride)
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()
@@ -487,4 +620,17 @@ class JsonlStore(Store, Streamable, Recordable):
if isinstance(keys, str):
return data[keys]
return {k: data[k] for k in keys}
return super().fetch_record(index, keys)
return _record_fetch(self, index, keys)
def fetch(self, begin: int, end: int, keys: Union[str, List[str]]):
if self._processor is not None:
raise RuntimeError(
"JsonlStore in lazy (processor) mode does not support "
"stream fetch(); use fetch_record() instead."
)
return _stream_fetch(self, begin, end, keys)
def __getitem__(self, index: int) -> Dict[str, Tensor]:
if self._processor is not None:
return self.fetch_record(index, self._record_keys())
return super().__getitem__(index)