From 663ef900fc04e0ddb65551dd3646655a0e40a6eb Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Sun, 19 Jul 2026 16:02:50 +0800 Subject: [PATCH] refactor: move sample-id indexing from dataset to store MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 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 --- astrai/dataset/dataset.py | 385 ++++++++++++++----------------------- astrai/dataset/storage.py | 368 ++++++++++++++++++++++++----------- tests/data/test_dataset.py | 116 ++++++----- 3 files changed, 468 insertions(+), 401 deletions(-) diff --git a/astrai/dataset/dataset.py b/astrai/dataset/dataset.py index 1f67ac7..58051bb 100644 --- a/astrai/dataset/dataset.py +++ b/astrai/dataset/dataset.py @@ -1,21 +1,26 @@ -"""Dataset implementations with factory pattern for training. +"""Dataset implementations for training. + +Composition over inheritance — every dataset is a thin wrapper that +binds a :class:`Store` to a particular train-type's key mapping. All +sample-id → token/record indexing lives on the Store; datasets never +know about window/stride math or segment layouts. 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 + BaseDataset (ABC) — holds a Store, exposes __len__/keys, + overrides __getitem__ + ├── SEQDataset — next-token prediction (stream) + ├── SFTDataset — loss-mask + position_ids (stream) + ├── DPODataset — chosen/rejected pairs (record) + └── GRPODataset — prompt + response group (record) -``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. +``DatasetFactory.load(train_type, load_path, window_size, stride, …)`` +builds the Store (auto-detecting format) before constructing the +matching dataset. Passing ``store=`` skips Store construction. -``__len__`` returns the sample count (stream: windows, record: -records) so DataLoader and progress bars work uniformly. +When a record dataset (DPO) reads from raw JSONL, a *processor* +function (pure ``record -> Dict[str, Tensor]``) is forwarded to +:class:`JsonlStore` so tokenisation happens on the fly. """ from abc import ABC, abstractmethod @@ -218,184 +223,57 @@ def grpo_collate_fn(batch: List[Dict[str, Tensor]]) -> Dict[str, Tensor]: } -class BaseDataset(Dataset, ABC): - """Abstract base class for all dataset types. +def validate_keys(store: Store, required: List[str]) -> None: + """Raise ``KeyError`` if *store* is missing any *required* key.""" + if not required: + return + actual = set(store.keys) + missing = [k for k in required if k not in actual] + if missing: + raise KeyError( + f"Store at {getattr(store, '_load_path', '?')} is missing required " + f"keys {missing}; available keys are {sorted(actual)}." + ) - Implements common functionality for window-based data fetching. - Uses a storage abstraction for format-agnostic data loading. + +class BaseDataset(Dataset, ABC): + """Abstract base class for dataset types. + + Holds a :class:`Store`. All sample-id indexing is delegated to the + store — this class exposes ``__len__`` as ``len(store)`` and the + ``keys`` property as ``store.keys``. Subclasses implement + ``__getitem__`` with the train-type-specific key mapping and any + training-only index arithmetic (e.g. the next-token ``+1`` shift). """ - def __init__(self, window_size: int, stride: int): + required_keys: List[str] = [] + + def __init__(self, store: Store): super().__init__() - self.window_size = window_size - self.stride = stride - self.storage: Optional[Store] = None + self.store: Store = store + validate_keys(store, self.required_keys) - @property - def required_keys(self) -> List[str]: - """Return required storage keys for this dataset type. - - Subclasses should override to specify expected keys. - """ - return [] - - def _validate_keys(self): - if not self.required_keys: - return - 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"Dataset {type(self).__name__} requires keys {self.required_keys}, " - f"but storage at {self._load_path} only has {sorted(actual_keys)}. " - f"Missing: {missing}" - ) - - def load(self, load_path: str, storage_type: Optional[str] = None, **kwargs): - """Load dataset from the given path. - - Auto-detects the storage format if not specified. - - Args: - load_path: Path to the data directory or file - storage_type: Force a specific storage type ("h5", "bin", "jsonl"), - or None for auto-detection - **kwargs: Extra arguments forwarded to the store constructor and - to ``store.load()``. - - Raises: - KeyError: If the loaded storage is missing required keys. - """ - 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() - - @property - def count(self) -> int: - """Return the total number of raw elements (tokens) in the dataset.""" - if self.storage is None: - return 0 - return len(self.storage) + def __len__(self) -> int: + return len(self.store) @property def keys(self) -> List[str]: - """Return the available data keys.""" - if self.storage is None: - return [] - return self.storage.keys + return self.store.keys - def get_index(self, index: int) -> tuple: - """Calculate begin and end indices for a sample. - - Args: - index: Sample index - - Returns: - Tuple of (begin_idx, end_idx) - """ - if self.storage is None: - raise RuntimeError("Dataset not loaded, call load() first") - total = len(self.storage) - if total <= self.window_size: - raise ValueError( - f"Data too short: {total} tokens <= window_size {self.window_size}" - ) - - begin_idx = min(index * self.stride, total - 1 - self.window_size) - end_idx = min(begin_idx + self.window_size, total - 1) - - return begin_idx, end_idx + @property + def token_count(self) -> int: + return self.store.token_count @abstractmethod def __getitem__(self, index: int) -> Dict[str, Tensor]: - """Get a single sample by index. - - Must be implemented by subclasses. - """ raise NotImplementedError - def __len__(self) -> int: - if self.storage is None: - return 0 - total = len(self.storage) - if total <= self.window_size: - return 0 - 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. + """Factory for creating dataset instances by train-type. - Supports decorator-based registration for extensible dataset types. - All default dataset types (seq, sft, dpo, grpo) are registered automatically - when their classes are defined with the decorator. - - Example usage: - @DatasetFactory.register("custom") - class CustomDataset(BaseDataset): - ... - - dataset = DatasetFactory.create("custom", window_size, stride) + Use :meth:`DatasetFactory.register("custom")` to register new + dataset classes; they must inherit from :class:`BaseDataset`. """ @classmethod @@ -417,32 +295,34 @@ class DatasetFactory(BaseFactory["BaseDataset"]): - **store given**: bind it directly — the caller fully controls Store construction and processor setup. *load_path*, - *storage_type*, *tokenizer_path* are ignored. + *storage_type*, *tokenizer_path*, *window_size*, *stride* are + ignored. - **store is None**: build a Store from *load_path*, auto-detecting format and constructing a processor when *tokenizer_path* is given for a record dataset on JSONL. Args: - train_type: Type of training dataset - load_path: Path to the data file (ignored if *store* given) - 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 for lazy JSONL tokenisation - max_len: Max sequence length for the processor - store: Pre-built, already-loaded Store instance - **kwargs: Extra arguments forwarded to ``dataset.load()`` + train_type: Registered dataset name ("seq", "sft", "dpo", + "grpo", …). + load_path: Path to the data file or directory (ignored if + *store* is given). + window_size: Stream window length — only meaningful for + stream datasets (SEQ/SFT). Record datasets ignore it. + stride: Stride between consecutive stream samples + (default: same as *window_size*). + storage_type: Storage backend ("h5", "bin", "jsonl") or + None for auto-detection. + tokenizer_path: Path to tokenizer for lazy JSONL + tokenisation (record datasets only). + max_len: Max sequence length forwarded to processors. + store: Pre-built, already-loaded Store instance. + **kwargs: Extra arguments forwarded to ``store.load()``. Returns: - Loaded dataset instance + Loaded dataset instance. """ - if stride is None: - stride = window_size - if store is not None: - dataset = cls.create(train_type, window_size, stride) - dataset.storage = store - return dataset + return cls.create(train_type, store=store) if load_path is None: raise ValueError("Either load_path or store must be provided") @@ -450,14 +330,37 @@ class DatasetFactory(BaseFactory["BaseDataset"]): if storage_type is None: storage_type = detect_format(load_path) + if stride is None: + stride = window_size + 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) + store_window = cls._store_window_for(train_type, window_size) + store = StoreFactory.create( + storage_type, + window_size=store_window, + stride=stride if stride else store_window, + ) + if processor is not None: + store.load(load_path, processor=processor, **kwargs) + else: + store.load(load_path, **kwargs) - return dataset + return cls.create(train_type, store=store) + + @staticmethod + def _store_window_for(train_type: str, window_size: int) -> int: + """Stream datasets consume ``window_size``; record datasets ignore it. + + Record datasets (dpo/grpo) treat each record as an independent + training unit and never window, so the store is built with + ``window_size=0`` and ``len(store)`` returns the record count. + """ + if train_type in ("seq", "sft"): + return window_size + return 0 @staticmethod def _maybe_build_processor( @@ -482,43 +385,41 @@ class DatasetFactory(BaseFactory["BaseDataset"]): @DatasetFactory.register("seq") class SEQDataset(BaseDataset): - """Dataset for sequential next-token prediction training.""" + """Dataset for sequential next-token prediction training. - @property - def required_keys(self) -> List[str]: - return ["sequence"] + Stream mode: ``store.fetch(begin, end, "sequence")`` returns the + input window; the +1 shifted call returns the next-token target. + """ - def _fetch_data(self, begin_idx: int, end_idx: int) -> Tensor: - return self.storage.fetch(begin_idx, end_idx, "sequence") + required_keys = ["sequence"] - def __getitem__(self, index): - begin_idx, end_idx = self.get_index(index) - - x = self._fetch_data(begin_idx, end_idx).to(dtype=torch.long) - y = self._fetch_data(begin_idx + 1, end_idx + 1).to(dtype=torch.long) - - return {"input_ids": x, "target_ids": y} + def __getitem__(self, index: int): + begin, end = self.store.sample_window(index) + x = self.store.fetch(begin, end, "sequence") + y = self.store.fetch(begin + 1, end + 1, "sequence") + return { + "input_ids": x.to(dtype=torch.long), + "target_ids": y.to(dtype=torch.long), + } @DatasetFactory.register("sft") class SFTDataset(BaseDataset): - """Dataset for supervised fine-tuning with loss masking.""" + """Dataset for supervised fine-tuning with loss masking. - @property - def required_keys(self) -> List[str]: - return ["sequence", "loss_mask", "position_ids"] + Stream mode: ``sequence``/``loss_mask``/``position_ids`` are sliced + to the window. ``loss_mask`` and ``target_ids`` use the +1 shifted + slice so they align with the predicted positions. + """ - def _fetch_data(self, begin_idx: int, end_idx: int, key: str) -> Tensor: - return self.storage.fetch(begin_idx, end_idx, key) - - def __getitem__(self, index): - begin_idx, end_idx = self.get_index(index) - - x = self._fetch_data(begin_idx, end_idx, "sequence") - y = self._fetch_data(begin_idx + 1, end_idx + 1, "sequence") - position_ids = self._fetch_data(begin_idx, end_idx, "position_ids") - loss_mask = self._fetch_data(begin_idx + 1, end_idx + 1, "loss_mask") + required_keys = ["sequence", "loss_mask", "position_ids"] + def __getitem__(self, index: int): + begin, end = self.store.sample_window(index) + x = self.store.fetch(begin, end, "sequence") + y = self.store.fetch(begin + 1, end + 1, "sequence") + position_ids = self.store.fetch(begin, end, "position_ids") + loss_mask = self.store.fetch(begin + 1, end + 1, "loss_mask") return { "input_ids": x.to(dtype=torch.long), "target_ids": y.to(dtype=torch.long), @@ -528,7 +429,7 @@ class SFTDataset(BaseDataset): @DatasetFactory.register("dpo") -class DPODataset(RecordDataset): +class DPODataset(BaseDataset): """Record-structured dataset for Direct Preference Optimization. Each sample is one preference pair (chosen + rejected) and is an @@ -536,39 +437,35 @@ class DPODataset(RecordDataset): concatenation. This keeps each sequence self-contained so attention never leaks across preference pairs. - Two loading paths (handled by :class:`RecordDataset`): + Two loading paths (handled by :class:`DatasetFactory`): - - **Pre-tokenized** (H5/bin): ``load(path)`` reads per-record tensors, - ``__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``. + - **Pre-tokenized** (H5/bin): ``store.load(path)`` reads per-record + tensors; ``__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"] + required_keys = ["chosen", "rejected", "chosen_mask", "rejected_mask"] def make_processor(self, tokenizer, max_len: int): return partial(dpo_processor, tokenizer=tokenizer, max_len=max_len) def __getitem__(self, index: int) -> Dict[str, Tensor]: return { - "chosen": self.storage.fetch_record(index, "chosen").to(dtype=torch.long), - "rejected": self.storage.fetch_record(index, "rejected").to( - dtype=torch.long - ), - "chosen_mask": self.storage.fetch_record(index, "chosen_mask").to( + "chosen": self.store.fetch_record(index, "chosen").to(dtype=torch.long), + "rejected": self.store.fetch_record(index, "rejected").to(dtype=torch.long), + "chosen_mask": self.store.fetch_record(index, "chosen_mask").to( dtype=torch.bool ), - "rejected_mask": self.storage.fetch_record(index, "rejected_mask").to( + "rejected_mask": self.store.fetch_record(index, "rejected_mask").to( dtype=torch.bool ), } @DatasetFactory.register("grpo") -class GRPODataset(RecordDataset): +class GRPODataset(BaseDataset): """Dataset for offline Group Relative Policy Optimization. Each sample is one prompt with its group of responses and scalar @@ -582,15 +479,13 @@ class GRPODataset(RecordDataset): - ``rewards``: List[Tensor] — one 1-D float tensor (len G) per record """ - @property - def required_keys(self) -> List[str]: - return ["prompts", "responses", "masks", "rewards"] + required_keys = ["prompts", "responses", "masks", "rewards"] def __getitem__(self, index: int) -> Dict[str, Tensor]: - prompts = self.storage.fetch_record(index, "prompts") - responses = self.storage.fetch_record(index, "responses") - masks = self.storage.fetch_record(index, "masks") - rewards = self.storage.fetch_record(index, "rewards") + prompts = self.store.fetch_record(index, "prompts") + responses = self.store.fetch_record(index, "responses") + masks = self.store.fetch_record(index, "masks") + rewards = self.store.fetch_record(index, "rewards") return { "prompts": prompts.to(dtype=torch.long), "responses": [r.to(dtype=torch.long) for r in responses], diff --git a/astrai/dataset/storage.py b/astrai/dataset/storage.py index c0d09a3..81bf5af 100644 --- a/astrai/dataset/storage.py +++ b/astrai/dataset/storage.py @@ -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) diff --git a/tests/data/test_dataset.py b/tests/data/test_dataset.py index fb4b599..29b12c2 100644 --- a/tests/data/test_dataset.py +++ b/tests/data/test_dataset.py @@ -7,7 +7,7 @@ import pytest import torch from astrai.config.preprocess_config import PipelineConfig -from astrai.dataset.dataset import DatasetFactory, SEQDataset, dpo_tokenize +from astrai.dataset.dataset import DatasetFactory, dpo_tokenize from astrai.dataset.storage import ( H5Store, JsonlStore, @@ -118,7 +118,7 @@ def test_dpo_strategy_with_random_data(base_test_env): ) assert dpo_dataset is not None - assert dpo_dataset.storage is not None + assert dpo_dataset.store is not None assert len(dpo_dataset) > 0 # Test that we can get DPO items without errors @@ -147,7 +147,7 @@ def test_sft_dataset_with_random_data(base_test_env): ) assert sft_dataset is not None - assert sft_dataset.storage is not None + assert sft_dataset.store is not None assert len(sft_dataset) > 0 # Test that we can get SFT items without errors @@ -178,39 +178,37 @@ def test_dataset_with_custom_stride(base_test_env): assert len(dataset) > len(default_stride_dataset) -def test_dataset_count_property(base_test_env): +def test_dataset_token_count_property(base_test_env): + """dataset.token_count exposes the raw stream token length.""" test_dir = base_test_env["test_dir"] dataset = _make_seq_dataset(test_dir, "count_test_data") - assert dataset.count == 200 - assert dataset.count > len(dataset) + assert dataset.token_count == 200 + assert dataset.token_count > len(dataset) assert len(dataset) == (200 - 1 - 64) // 64 + 1 -def test_empty_dataset_count(): - """Test count returns 0 when no data is loaded""" - dataset = SEQDataset(window_size=64, stride=32) - assert dataset.count == 0 - assert dataset.keys == [] - - def test_dataset_too_short_for_window(base_test_env): test_dir = base_test_env["test_dir"] dataset = _make_seq_dataset(test_dir, "short", seq_length=30) assert len(dataset) == 0 - assert dataset.count == 30 + assert dataset.token_count == 30 -def test_unloaded_dataset_getitem_raises(): - """__getitem__ without load() should fail clearly""" - dataset = SEQDataset(window_size=64, stride=32) - with pytest.raises(RuntimeError, match="not loaded"): - dataset.get_index(0) +def test_unloaded_sample_window_raises(): + """Store.sample_window before load raises RuntimeError.""" + from astrai.dataset.storage import H5Store + + store = H5Store(window_size=64, stride=64) + with pytest.raises(IndexError, match="Data too short"): + store.sample_window(0) def test_unloaded_dataset_len(): - """__len__ without load() returns 0""" - dataset = SEQDataset(window_size=64, stride=32) - assert len(dataset) == 0 + """__len__ on a store with no data returns 0.""" + from astrai.dataset.storage import H5Store + + store = H5Store(window_size=64, stride=64) + assert len(store) == 0 def test_store_unloaded_len(): @@ -223,7 +221,7 @@ def test_store_unloaded_len(): def test_store_fetch_begin_equals_end(base_test_env): test_dir = base_test_env["test_dir"] dataset = _make_seq_dataset(test_dir, "empty_fetch", seq_length=100, window_size=32) - result = dataset.storage.fetch(10, 10, "sequence") + result = dataset.store.fetch(10, 10, "sequence") assert result.numel() == 0 @@ -273,7 +271,7 @@ def test_store_multi_segment_concat(base_test_env): store = StoreFactory.create("h5") store.load(data_dir) - assert len(store) == 9 + assert store.token_count == 9 result = store.fetch(2, 7, "sequence") assert result.tolist() == [3, 4, 5, 6, 7] @@ -302,7 +300,9 @@ def test_mmap_store_load_and_fetch(base_test_env): store = StoreFactory.create("bin") store.load(test_dir) - assert len(store) == 200 + assert store.token_count == 200 + assert store.num_records == 0 + assert len(store) == 0 # no window configured, no records → 0 samples assert "sequence" in store.keys result = store.fetch(10, 20, "sequence") @@ -315,23 +315,26 @@ def test_mmap_dataset_load(base_test_env): save_bin(test_dir, data) dataset = DatasetFactory.load("seq", test_dir, window_size=64) assert len(dataset) > 0 - assert dataset.count == 200 + assert dataset.token_count == 200 assert dataset[0]["input_ids"].shape[0] == 64 def test_normalize_empty_key(): - """_normalize with empty tensor list does not crash""" + """_normalize with empty tensor list does not crash.""" store = H5Store() store._normalize({"sequence": []}) assert len(store) == 0 + assert store.num_records == 0 # empty key forces num_records=0 assert store.keys == ["sequence"] def test_normalize_mixed_empty_key(): - """_normalize with empty + non-empty keys returns min=0""" + """_normalize with empty + non-empty keys returns min=0 records.""" store = H5Store() store._normalize({"sequence": [torch.tensor([1, 2, 3])], "loss_mask": []}) assert len(store) == 0 + assert store.num_records == 0 + assert store.token_count == 0 # min() over keys assert set(store.keys) == {"sequence", "loss_mask"} @@ -339,15 +342,14 @@ def test_grpo_dataset_dtype(base_test_env): """GRPO dataset returns correct dtypes for per-record structured data.""" from astrai.dataset.dataset import GRPODataset - test_dir = base_test_env["test_dir"] G = 4 - dataset = GRPODataset() - dataset.storage = type( + store = type( "FakeStore", (), { "keys": ["prompts", "responses", "masks", "rewards"], "num_records": 1, + "token_count": 0, "_data": { "prompts": [torch.randint(0, 100, (10,), dtype=torch.int32)], "responses": [ @@ -357,8 +359,10 @@ def test_grpo_dataset_dtype(base_test_env): "rewards": [torch.rand(G, dtype=torch.float32)], }, "fetch_record": _fake_fetch_record, + "__len__": lambda self: self.num_records, }, )() + dataset = GRPODataset(store=store) item = dataset[0] assert item["prompts"].dtype == torch.long @@ -371,17 +375,16 @@ def test_grpo_dataset_load(base_test_env): """GRPO dataset loads record-structured data with per-response boundaries.""" from astrai.dataset.dataset import GRPODataset - test_dir = base_test_env["test_dir"] G = 3 prompt_len = 8 resp_lens = [5, 7, 4] - dataset = GRPODataset() - dataset.storage = type( + store = type( "FakeStore", (), { "keys": ["prompts", "responses", "masks", "rewards"], "num_records": 1, + "token_count": 0, "_data": { "prompts": [torch.randint(0, 100, (prompt_len,))], "responses": [[torch.randint(0, 100, (rl,)) for rl in resp_lens]], @@ -389,8 +392,10 @@ def test_grpo_dataset_load(base_test_env): "rewards": [torch.tensor([0.9, 0.3, 0.7], dtype=torch.float32)], }, "fetch_record": _fake_fetch_record, + "__len__": lambda self: self.num_records, }, )() + dataset = GRPODataset(store=store) assert len(dataset) == 1 item = dataset[0] @@ -458,7 +463,7 @@ def test_dataset_load_explicit_storage_type(base_test_env): test_dir = base_test_env["test_dir"] dataset = _make_seq_dataset(test_dir, "explicit", storage_type="h5") assert len(dataset) > 0 - assert dataset.count == 200 + assert dataset.token_count == 200 def _write_json_dataset(test_dir, tokenizer_path, records, config_overrides=None): @@ -853,11 +858,11 @@ def test_grpo_collate_variable_lengths(): assert result["responses"][0, 0, 0] == 4 assert result["responses"][0, 0, 1] == 5 assert result["responses"][0, 0, 2] == 0 # padded - assert result["masks"][0, 0, 2] == False # padded + assert not result["masks"][0, 0, 2] # padded # Check response content: item 0, response 1 is [6,7,8,9] no padding assert result["responses"][0, 1, 3] == 9 - assert result["masks"][0, 1, 3] == True + assert result["masks"][0, 1, 3] def test_grpo_multiple_records(base_test_env): @@ -871,13 +876,13 @@ def test_grpo_multiple_records(base_test_env): [torch.randint(0, 100, (np.random.randint(3, 8),)) for _ in range(G)] for _ in range(n_records) ] - dataset = GRPODataset() - dataset.storage = type( + store = type( "FakeStore", (), { "keys": ["prompts", "responses", "masks", "rewards"], "num_records": n_records, + "token_count": 0, "_data": { "prompts": [torch.randint(0, 100, (10,)) for _ in range(n_records)], "responses": dummy_responses, @@ -890,8 +895,10 @@ def test_grpo_multiple_records(base_test_env): ], }, "fetch_record": _fake_fetch_record, + "__len__": lambda self: self.num_records, }, )() + dataset = GRPODataset(store=store) assert len(dataset) == n_records @@ -973,8 +980,8 @@ def test_dpo_jsonl_lazy_load(base_test_env): ) assert len(ds) == 2 - assert ds.storage.num_records == 2 - assert ds.storage._processor is not None + assert ds.store.num_records == 2 + assert ds.store._processor is not None item = ds[0] assert set(item.keys()) == {"chosen", "rejected", "chosen_mask", "rejected_mask"} @@ -1035,7 +1042,13 @@ def test_jsonl_store_eager_len_returns_token_count(base_test_env): def test_h5_store_dual_mode(base_test_env): - """H5Store supports both fetch (stream) and fetch_record (record).""" + """H5Store supports both fetch (stream) and fetch_record (record). + + No window configured → ``len(store)`` reflects the record count + (2). ``token_count`` retains the legacy stream length (128), and + token-stream access via :meth:`fetch` is still available for + callers that want explicit begin/end control. + """ test_dir = base_test_env["test_dir"] seq_length = 64 @@ -1048,8 +1061,9 @@ def test_h5_store_dual_mode(base_test_env): store = H5Store() store.load(test_dir) - assert len(store) == seq_length * 2 + assert store.token_count == seq_length * 2 assert store.num_records == 2 + assert len(store) == 2 # no window configured → record count rec0 = store.fetch_record(0, "chosen") assert rec0.shape == (seq_length,) @@ -1057,9 +1071,20 @@ def test_h5_store_dual_mode(base_test_env): stream = store.fetch(0, 10, "chosen") assert stream.shape == (10,) + # Window-configured view of the same data uses stream sample count: + # token_count=128, window_size=64 → num_samples = (128-1-64)//64 + 1 = 1 + stream_view = H5Store(window_size=seq_length, stride=seq_length) + stream_view.load(test_dir) + assert len(stream_view) == 1 + def test_mmap_store_stream_only_no_offsets(base_test_env): - """MmapStore without offsets: num_records == 0, stream works.""" + """MmapStore without offsets: num_records == 0, stream works. + + No window configured → ``len(store)`` is 0 (no iterate units). + ``token_count`` remains 128 for raw token slicing, and ``fetch`` + provides direct token-range access. + """ test_dir = base_test_env["test_dir"] seq_length = 128 @@ -1069,8 +1094,9 @@ def test_mmap_store_stream_only_no_offsets(base_test_env): store = StoreFactory.create("bin") store.load(test_dir) - assert len(store) == seq_length + assert store.token_count == seq_length assert store.num_records == 0 + assert len(store) == 0 chunk = store.fetch(0, 32, "sequence") assert chunk.shape == (32,)