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/dataset.py
+140 -245
View File
@@ -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: Class hierarchy:
BaseDataset (ABC) — load/validate, owns a Store BaseDataset (ABC) — holds a Store, exposes __len__/keys,
├── SEQDataset — stream, next-token prediction (PT) overrides __getitem__
├── SFTDataset — stream, loss-mask + position_ids ├── SEQDataset — next-token prediction (stream)
└── RecordDataset — record access, optional processor ├── SFTDataset — loss-mask + position_ids (stream)
├── DPODataset — chosen/rejected pairs ├── DPODataset — chosen/rejected pairs (record)
└── GRPODataset — prompt + response group └── GRPODataset — prompt + response group (record)
``RecordDataset`` holds an optional *processor* (pure ``DatasetFactory.load(train_type, load_path, window_size, stride, …)``
``record -> Dict[str, Tensor]`` function). When the backing Store is builds the Store (auto-detecting format) before constructing the
a lazy JsonlStore, the processor tokenises on the fly; otherwise it matching dataset. Passing ``store=`` skips Store construction.
is ignored and ``fetch_record`` reads pre-tokenised tensors.
``__len__`` returns the sample count (stream: windows, record: When a record dataset (DPO) reads from raw JSONL, a *processor*
records) so DataLoader and progress bars work uniformly. function (pure ``record -> Dict[str, Tensor]``) is forwarded to
:class:`JsonlStore` so tokenisation happens on the fly.
""" """
from abc import ABC, abstractmethod 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): def validate_keys(store: Store, required: List[str]) -> None:
"""Abstract base class for all dataset types. """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__() super().__init__()
self.window_size = window_size self.store: Store = store
self.stride = stride validate_keys(store, self.required_keys)
self.storage: Optional[Store] = None
@property def __len__(self) -> int:
def required_keys(self) -> List[str]: return len(self.store)
"""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)
@property @property
def keys(self) -> List[str]: def keys(self) -> List[str]:
"""Return the available data keys.""" return self.store.keys
if self.storage is None:
return []
return self.storage.keys
def get_index(self, index: int) -> tuple: @property
"""Calculate begin and end indices for a sample. def token_count(self) -> int:
return self.store.token_count
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
@abstractmethod @abstractmethod
def __getitem__(self, index: int) -> Dict[str, Tensor]: def __getitem__(self, index: int) -> Dict[str, Tensor]:
"""Get a single sample by index.
Must be implemented by subclasses.
"""
raise NotImplementedError 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"]): 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. Use :meth:`DatasetFactory.register("custom")` to register new
All default dataset types (seq, sft, dpo, grpo) are registered automatically dataset classes; they must inherit from :class:`BaseDataset`.
when their classes are defined with the decorator.
Example usage:
@DatasetFactory.register("custom")
class CustomDataset(BaseDataset):
...
dataset = DatasetFactory.create("custom", window_size, stride)
""" """
@classmethod @classmethod
@@ -417,32 +295,34 @@ class DatasetFactory(BaseFactory["BaseDataset"]):
- **store given**: bind it directly — the caller fully controls - **store given**: bind it directly — the caller fully controls
Store construction and processor setup. *load_path*, 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 - **store is None**: build a Store from *load_path*, auto-detecting
format and constructing a processor when *tokenizer_path* is format and constructing a processor when *tokenizer_path* is
given for a record dataset on JSONL. given for a record dataset on JSONL.
Args: Args:
train_type: Type of training dataset train_type: Registered dataset name ("seq", "sft", "dpo",
load_path: Path to the data file (ignored if *store* given) "grpo", …).
window_size: Window size for data sampling load_path: Path to the data file or directory (ignored if
stride: Stride between consecutive samples (default: same as window_size) *store* is given).
storage_type: Storage type ("h5", "bin", "jsonl") or None for auto-detection window_size: Stream window length — only meaningful for
tokenizer_path: Path to tokenizer for lazy JSONL tokenisation stream datasets (SEQ/SFT). Record datasets ignore it.
max_len: Max sequence length for the processor stride: Stride between consecutive stream samples
store: Pre-built, already-loaded Store instance (default: same as *window_size*).
**kwargs: Extra arguments forwarded to ``dataset.load()`` 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: Returns:
Loaded dataset instance Loaded dataset instance.
""" """
if stride is None:
stride = window_size
if store is not None: if store is not None:
dataset = cls.create(train_type, window_size, stride) return cls.create(train_type, store=store)
dataset.storage = store
return dataset
if load_path is None: if load_path is None:
raise ValueError("Either load_path or store must be provided") raise ValueError("Either load_path or store must be provided")
@@ -450,14 +330,37 @@ class DatasetFactory(BaseFactory["BaseDataset"]):
if storage_type is None: if storage_type is None:
storage_type = detect_format(load_path) storage_type = detect_format(load_path)
if stride is None:
stride = window_size
processor = cls._maybe_build_processor( processor = cls._maybe_build_processor(
train_type, storage_type, tokenizer_path, max_len train_type, storage_type, tokenizer_path, max_len
) )
dataset = cls.create(train_type, window_size, stride, processor=processor) store_window = cls._store_window_for(train_type, window_size)
dataset.load(load_path, storage_type=storage_type, **kwargs) 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 @staticmethod
def _maybe_build_processor( def _maybe_build_processor(
@@ -482,43 +385,41 @@ class DatasetFactory(BaseFactory["BaseDataset"]):
@DatasetFactory.register("seq") @DatasetFactory.register("seq")
class SEQDataset(BaseDataset): class SEQDataset(BaseDataset):
"""Dataset for sequential next-token prediction training.""" """Dataset for sequential next-token prediction training.
@property Stream mode: ``store.fetch(begin, end, "sequence")`` returns the
def required_keys(self) -> List[str]: input window; the +1 shifted call returns the next-token target.
return ["sequence"] """
def _fetch_data(self, begin_idx: int, end_idx: int) -> Tensor: required_keys = ["sequence"]
return self.storage.fetch(begin_idx, end_idx, "sequence")
def __getitem__(self, index): def __getitem__(self, index: int):
begin_idx, end_idx = self.get_index(index) begin, end = self.store.sample_window(index)
x = self.store.fetch(begin, end, "sequence")
x = self._fetch_data(begin_idx, end_idx).to(dtype=torch.long) y = self.store.fetch(begin + 1, end + 1, "sequence")
y = self._fetch_data(begin_idx + 1, end_idx + 1).to(dtype=torch.long) return {
"input_ids": x.to(dtype=torch.long),
return {"input_ids": x, "target_ids": y} "target_ids": y.to(dtype=torch.long),
}
@DatasetFactory.register("sft") @DatasetFactory.register("sft")
class SFTDataset(BaseDataset): class SFTDataset(BaseDataset):
"""Dataset for supervised fine-tuning with loss masking.""" """Dataset for supervised fine-tuning with loss masking.
@property Stream mode: ``sequence``/``loss_mask``/``position_ids`` are sliced
def required_keys(self) -> List[str]: to the window. ``loss_mask`` and ``target_ids`` use the +1 shifted
return ["sequence", "loss_mask", "position_ids"] slice so they align with the predicted positions.
"""
def _fetch_data(self, begin_idx: int, end_idx: int, key: str) -> Tensor: required_keys = ["sequence", "loss_mask", "position_ids"]
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")
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 { return {
"input_ids": x.to(dtype=torch.long), "input_ids": x.to(dtype=torch.long),
"target_ids": y.to(dtype=torch.long), "target_ids": y.to(dtype=torch.long),
@@ -528,7 +429,7 @@ class SFTDataset(BaseDataset):
@DatasetFactory.register("dpo") @DatasetFactory.register("dpo")
class DPODataset(RecordDataset): class DPODataset(BaseDataset):
"""Record-structured dataset for Direct Preference Optimization. """Record-structured dataset for Direct Preference Optimization.
Each sample is one preference pair (chosen + rejected) and is an 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 concatenation. This keeps each sequence self-contained so attention
never leaks across preference pairs. 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, - **Pre-tokenized** (H5/bin): ``store.load(path)`` reads per-record
``__getitem__`` returns them directly. tensors; ``__getitem__`` returns them directly.
- **Raw JSONL** (``tokenizer_path=...``): builds a lazy processor via - **Raw JSONL** (``tokenizer_path=...``): builds a lazy processor
:func:`dpo_processor` that tokenises on the fly — no packing, no via :func:`dpo_processor` that tokenises on the fly — no packing,
``position_ids``. no ``position_ids``.
""" """
@property required_keys = ["chosen", "rejected", "chosen_mask", "rejected_mask"]
def required_keys(self) -> List[str]:
return ["chosen", "rejected", "chosen_mask", "rejected_mask"]
def make_processor(self, tokenizer, max_len: int): def make_processor(self, tokenizer, max_len: int):
return partial(dpo_processor, tokenizer=tokenizer, max_len=max_len) return partial(dpo_processor, tokenizer=tokenizer, max_len=max_len)
def __getitem__(self, index: int) -> Dict[str, Tensor]: def __getitem__(self, index: int) -> Dict[str, Tensor]:
return { return {
"chosen": self.storage.fetch_record(index, "chosen").to(dtype=torch.long), "chosen": self.store.fetch_record(index, "chosen").to(dtype=torch.long),
"rejected": self.storage.fetch_record(index, "rejected").to( "rejected": self.store.fetch_record(index, "rejected").to(dtype=torch.long),
dtype=torch.long "chosen_mask": self.store.fetch_record(index, "chosen_mask").to(
),
"chosen_mask": self.storage.fetch_record(index, "chosen_mask").to(
dtype=torch.bool 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 dtype=torch.bool
), ),
} }
@DatasetFactory.register("grpo") @DatasetFactory.register("grpo")
class GRPODataset(RecordDataset): class GRPODataset(BaseDataset):
"""Dataset for offline Group Relative Policy Optimization. """Dataset for offline Group Relative Policy Optimization.
Each sample is one prompt with its group of responses and scalar 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 - ``rewards``: List[Tensor] — one 1-D float tensor (len G) per record
""" """
@property required_keys = ["prompts", "responses", "masks", "rewards"]
def required_keys(self) -> List[str]:
return ["prompts", "responses", "masks", "rewards"]
def __getitem__(self, index: int) -> Dict[str, Tensor]: def __getitem__(self, index: int) -> Dict[str, Tensor]:
prompts = self.storage.fetch_record(index, "prompts") prompts = self.store.fetch_record(index, "prompts")
responses = self.storage.fetch_record(index, "responses") responses = self.store.fetch_record(index, "responses")
masks = self.storage.fetch_record(index, "masks") masks = self.store.fetch_record(index, "masks")
rewards = self.storage.fetch_record(index, "rewards") rewards = self.store.fetch_record(index, "rewards")
return { return {
"prompts": prompts.to(dtype=torch.long), "prompts": prompts.to(dtype=torch.long),
"responses": [r.to(dtype=torch.long) for r in responses], "responses": [r.to(dtype=torch.long) for r in responses],
astrai/dataset/storage.py
+257 -111
View File
@@ -1,11 +1,14 @@
"""Storage backends for different data formats. """Storage backends for different data formats.
Architecture (mixin composition, no diamond inheritance): Architecture (composition over inheritance):
Store (ABC) — shared _data/_cum/_offsets bookkeeping Store (ABC) — owns _data/_cum/_offsets bookkeeping
+ _normalize() for registering segments + window_size/stride for sample-id
Streamable (mixin) — fetch(begin, end, key) for stream access indexing. __getitem__/__len__ produce
Recordable (mixin) — fetch_record(i, key) for record access 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) H5Store(Store, Streamable, Recordable)
MmapStore(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. Each mixin is a stateless trait that relies on ``self._data`` etc.
provided by :class:`Store`. Concrete stores mix in whichever access provided by :class:`Store`. Concrete stores mix in whichever access
modes they support — ``Store`` is the sole base class, so there is no primitives they support — ``Store`` is the sole base class, so there is
diamond inheritance or MRO ambiguity. 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 - **Stream mode** (``window_size > 0``): ``len(store)`` returns the number
concatenated segments. ``len(store)`` returns the total token count. of ``(window_size, stride)`` windows that fit in the token river;
- **Record** (DPO/GRPO): ``fetch_record(i, key)`` returns the *i*-th ``store[i]`` returns the *i*-th window as a dict of per-key tensors;
record without cross-record concatenation. ``num_records`` returns ``store.sample_window(i)`` exposes the underlying ``(begin, end)``
the record count. 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) ``segments_are_records`` (class attribute on each Store subclass)
tells ``_normalize`` whether segments are inherently per-record (H5/ tells ``_normalize`` whether segments are inherently per-record (H5/
@@ -40,7 +51,7 @@ import json
import logging import logging
from abc import ABC, abstractmethod from abc import ABC, abstractmethod
from pathlib import Path from pathlib import Path
from typing import Callable, Dict, List, Optional, Union from typing import Callable, Dict, List, Optional, Tuple, Union
import torch import torch
from torch import Tensor from torch import Tensor
@@ -102,20 +113,46 @@ def detect_format(load_path: str) -> str:
class Store(ABC): class Store(ABC):
"""Common base for all storage backends. """Common base for all storage backends.
Owns the shared ``_data`` / ``_cum`` / ``_offsets`` bookkeeping and A Store owns both its data layout AND its sample-id → token/record
the ``_normalize`` entry point used by tensor-backed subclasses. index translation. Datasets are thin wrappers that bind a Store
Does **not** expose an access API — that is the job of to a particular train-type's key mapping; they never know about
:class:`StreamStore` and :class:`RecordStore`. 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 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._data: Dict[str, List[Tensor]] = {}
self._cum: Dict[str, List[int]] = {} self._cum: Dict[str, List[int]] = {}
self._offsets: Dict[str, List[int]] = {} self._offsets: Dict[str, List[int]] = {}
self._length: int = 0 self._length: int = 0
self._num_records: 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 @abstractmethod
def load(self, path: str, **kwargs) -> None: def load(self, path: str, **kwargs) -> None:
@@ -125,14 +162,98 @@ class Store(ABC):
def keys(self) -> List[str]: def keys(self) -> List[str]:
return list(self._data.keys()) return list(self._data.keys())
def __len__(self) -> int: @property
"""Default: token count (stream semantics). def window_size(self) -> int:
return self._window_size
Subclasses that are record-only (e.g. lazy JsonlStore) override @property
to return ``self._num_records``. 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 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( def _normalize(
self, self,
raw: Dict[str, list], raw: Dict[str, list],
@@ -141,14 +262,14 @@ class Store(ABC):
"""Register segments and pre-compute indices for both access modes. """Register segments and pre-compute indices for both access modes.
Stream mode: ``_cum[key]`` accumulates per-segment lengths so Stream mode: ``_cum[key]`` accumulates per-segment lengths so
``StreamStore._fetch_key`` can bisect across segments without ``Streamable._fetch_stream_key`` can bisect across segments
concatenation. without concatenation.
Record mode: if *offsets* is provided (bin layout), Record mode: if *offsets* is provided (bin layout),
``_offsets[key]`` stores cumulative per-record offsets into the ``_offsets[key]`` stores cumulative per-record offsets into the
single concatenated segment. Otherwise, when single concatenated segment. Otherwise, when
``segments_are_records`` is True (H5/JSONL), ``_data[key]`` is a ``segments_are_records`` is True (H5/JSONL), ``_data[key]`` is
per-record list and ``fetch_record`` indexes it directly. a per-record list and ``fetch_record`` indexes it directly.
Nested keys (GRPO ``responses``/``masks`` as Nested keys (GRPO ``responses``/``masks`` as
``List[List[Tensor]]``) are stored as-is and excluded from both ``List[List[Tensor]]``) are stored as-is and excluded from both
@@ -194,7 +315,7 @@ class Store(ABC):
elif self.segments_are_records: elif self.segments_are_records:
per_record_counts = [] per_record_counts = []
for key, tensors in self._data.items(): for key, tensors in self._data.items():
if not tensors or isinstance(tensors[0], list): if tensors and isinstance(tensors[0], list):
continue continue
per_record_counts.append(len(tensors)) per_record_counts.append(len(tensors))
self._num_records = min(per_record_counts) if per_record_counts else 0 self._num_records = min(per_record_counts) if per_record_counts else 0
@@ -203,11 +324,13 @@ class Store(ABC):
class Streamable: 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``, Stateless trait relying on ``self._data``, ``self._cum``,
``self._length`` provided by :class:`Store`. Used by SEQ/SFT ``self._length`` maintained by :class:`Store`. Stream mode is
where data is a long token stream. 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( def fetch(
@@ -216,72 +339,75 @@ class Streamable:
end: int, end: int,
keys: Union[str, List[str]], keys: Union[str, List[str]],
): ):
if not self._data: return _stream_fetch(self, begin, end, keys)
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}
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 = [] def _stream_fetch(self, begin: int, end: int, keys: Union[str, List[str]]):
for i in range(seg_start, seg_end + 1): if not getattr(self, "_data", None):
prev = cum[i - 1] if i > 0 else 0 raise RuntimeError("Store not loaded")
s = max(begin - prev, 0) if not (0 <= begin < self._length and 0 <= end <= self._length):
e = min(end - prev, segments[i].shape[0]) raise ValueError(
results.append(segments[i][s:e]) 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: 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``, Stateless trait relying on ``self._data``, ``self._offsets``,
``self._num_records`` provided by :class:`Store`. Used by ``self._num_records`` maintained by :class:`Store`.
DPO/GRPO where each record is an independent training unit.
""" """
segments_are_records = True
@property
def num_records(self) -> int:
return self._num_records
def fetch_record( def fetch_record(
self, self,
index: int, index: int,
keys: Union[str, List[str]], keys: Union[str, List[str]],
): ):
if not self._data and self._num_records == 0: return _record_fetch(self, index, keys)
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) def _record_fetch(self, index: int, keys: Union[str, List[str]]):
if offsets: if not getattr(self, "_data", None) and self._num_records == 0:
start = offsets[index] raise RuntimeError("Store not loaded")
end = ( if not 0 <= index < self._num_records:
offsets[index + 1] raise ValueError(
if index + 1 < len(offsets) f"Record index out of bounds: {index}, num_records={self._num_records}"
else self._data[key][0].shape[0] )
) if isinstance(keys, str):
return self._data[key][0][start:end] return _fetch_record_key(self, keys, index)
return self._data[key][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"]): 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``, Each key is stored as a group of per-record datasets (``data_0``,
``data_1``, …). Supports both access modes: ``data_1``, …). Supports both access modes:
- **Stream** (``fetch(begin, end, key)``): concatenates across - **Stream**: ``fetch(begin, end, key)`` and ``store[i]`` slice
records via ``_cum`` — used by SEQ/SFT where data is a token stream. across concatenated records via ``_cum`` — used by SEQ/SFT.
- **Record** (``fetch_record(i, key)``): indexes ``_data[key]`` - **Record**: ``fetch_record(i, key)`` and ``store[i]`` (when
directly — used by DPO/GRPO where each record is independent. ``window_size == 0``) index ``_data[key]`` directly — used by
DPO/GRPO.
``len(store)`` returns the **token count** (stream semantics) so
SEQ/SFT windowing works. Record-only code uses
``store.num_records`` instead.
""" """
segments_are_records = True 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): def load(self, path: str, **kwargs):
self._normalize(load_h5(path)) self._normalize(load_h5(path))
@@ -321,11 +451,12 @@ class MmapStore(Store, Streamable, Recordable):
Supports both access modes: 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`` - **Record** (``fetch_record(i, key)``): only when ``meta.json``
contains per-record ``offsets`` (written via contains per-record ``offsets`` (written via
``save_bin(..., record_keys=...)``). Legacy bin files without ``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 ``segments_are_records`` is ``False`` here (bin segments are
contiguous streams, not per-record) — record access is driven contiguous streams, not per-record) — record access is driven
@@ -334,6 +465,14 @@ class MmapStore(Store, Streamable, Recordable):
segments_are_records = False 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): def load(self, path: str, **kwargs):
self._mmap_refs = [] self._mmap_refs = []
root = Path(path) root = Path(path)
@@ -416,32 +555,26 @@ class JsonlStore(Store, Streamable, Recordable):
- **Eager** (default): applies a :class:`TokenizeTransform` to every - **Eager** (default): applies a :class:`TokenizeTransform` to every
record at load time and registers per-key tensors via record at load time and registers per-key tensors via
``_normalize``. Both stream (``fetch``) and record ``_normalize``. Both ``fetch`` (stream) and ``fetch_record``
(``fetch_record``) access work — stream concatenates across (record) work.
records via ``_cum``, record indexes directly. - **Lazy** (``processor=fn`` passed): keeps raw records and defers
- **Lazy** (``processor=fn`` given): keeps raw records and defers tokenisation to ``fetch_record``. Only record access works —
tokenisation to ``fetch_record``. Only record access works. ``len(store)`` returns ``num_records``; stream primitives raise.
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" CONFIG_NAME = "dataset_config.json"
segments_are_records = True segments_are_records = True
def __init__(self): def __init__(
super().__init__() self,
window_size: int = 0,
stride: Optional[int] = None,
):
super().__init__(window_size=window_size, stride=stride)
self._source: Optional[JsonlSource] = None self._source: Optional[JsonlSource] = None
self._processor: Optional[Callable[[dict], Dict[str, Tensor]]] = None self._processor: Optional[Callable[[dict], Dict[str, Tensor]]] = None
self._keys_cache: Optional[List[str]] = 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): def load(self, path: str, transform=None, processor=None, **kwargs):
self._source = JsonlSource(path) self._source = JsonlSource(path)
records = self._source.load() records = self._source.load()
@@ -487,4 +620,17 @@ class JsonlStore(Store, Streamable, Recordable):
if isinstance(keys, str): if isinstance(keys, str):
return data[keys] return data[keys]
return {k: data[k] for k in 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)
tests/data/test_dataset.py
+71 -45
View File
@@ -7,7 +7,7 @@ import pytest
import torch import torch
from astrai.config.preprocess_config import PipelineConfig 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 ( from astrai.dataset.storage import (
H5Store, H5Store,
JsonlStore, JsonlStore,
@@ -118,7 +118,7 @@ def test_dpo_strategy_with_random_data(base_test_env):
) )
assert dpo_dataset is not None 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 assert len(dpo_dataset) > 0
# Test that we can get DPO items without errors # 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 is not None
assert sft_dataset.storage is not None assert sft_dataset.store is not None
assert len(sft_dataset) > 0 assert len(sft_dataset) > 0
# Test that we can get SFT items without errors # 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) 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"] test_dir = base_test_env["test_dir"]
dataset = _make_seq_dataset(test_dir, "count_test_data") dataset = _make_seq_dataset(test_dir, "count_test_data")
assert dataset.count == 200 assert dataset.token_count == 200
assert dataset.count > len(dataset) assert dataset.token_count > len(dataset)
assert len(dataset) == (200 - 1 - 64) // 64 + 1 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): def test_dataset_too_short_for_window(base_test_env):
test_dir = base_test_env["test_dir"] test_dir = base_test_env["test_dir"]
dataset = _make_seq_dataset(test_dir, "short", seq_length=30) dataset = _make_seq_dataset(test_dir, "short", seq_length=30)
assert len(dataset) == 0 assert len(dataset) == 0
assert dataset.count == 30 assert dataset.token_count == 30
def test_unloaded_dataset_getitem_raises(): def test_unloaded_sample_window_raises():
"""__getitem__ without load() should fail clearly""" """Store.sample_window before load raises RuntimeError."""
dataset = SEQDataset(window_size=64, stride=32) from astrai.dataset.storage import H5Store
with pytest.raises(RuntimeError, match="not loaded"):
dataset.get_index(0) store = H5Store(window_size=64, stride=64)
with pytest.raises(IndexError, match="Data too short"):
store.sample_window(0)
def test_unloaded_dataset_len(): def test_unloaded_dataset_len():
"""__len__ without load() returns 0""" """__len__ on a store with no data returns 0."""
dataset = SEQDataset(window_size=64, stride=32) from astrai.dataset.storage import H5Store
assert len(dataset) == 0
store = H5Store(window_size=64, stride=64)
assert len(store) == 0
def test_store_unloaded_len(): def test_store_unloaded_len():
@@ -223,7 +221,7 @@ def test_store_unloaded_len():
def test_store_fetch_begin_equals_end(base_test_env): def test_store_fetch_begin_equals_end(base_test_env):
test_dir = base_test_env["test_dir"] test_dir = base_test_env["test_dir"]
dataset = _make_seq_dataset(test_dir, "empty_fetch", seq_length=100, window_size=32) 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 assert result.numel() == 0
@@ -273,7 +271,7 @@ def test_store_multi_segment_concat(base_test_env):
store = StoreFactory.create("h5") store = StoreFactory.create("h5")
store.load(data_dir) store.load(data_dir)
assert len(store) == 9 assert store.token_count == 9
result = store.fetch(2, 7, "sequence") result = store.fetch(2, 7, "sequence")
assert result.tolist() == [3, 4, 5, 6, 7] 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 = StoreFactory.create("bin")
store.load(test_dir) 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 assert "sequence" in store.keys
result = store.fetch(10, 20, "sequence") result = store.fetch(10, 20, "sequence")
@@ -315,23 +315,26 @@ def test_mmap_dataset_load(base_test_env):
save_bin(test_dir, data) save_bin(test_dir, data)
dataset = DatasetFactory.load("seq", test_dir, window_size=64) dataset = DatasetFactory.load("seq", test_dir, window_size=64)
assert len(dataset) > 0 assert len(dataset) > 0
assert dataset.count == 200 assert dataset.token_count == 200
assert dataset[0]["input_ids"].shape[0] == 64 assert dataset[0]["input_ids"].shape[0] == 64
def test_normalize_empty_key(): def test_normalize_empty_key():
"""_normalize with empty tensor list does not crash""" """_normalize with empty tensor list does not crash."""
store = H5Store() store = H5Store()
store._normalize({"sequence": []}) store._normalize({"sequence": []})
assert len(store) == 0 assert len(store) == 0
assert store.num_records == 0 # empty key forces num_records=0
assert store.keys == ["sequence"] assert store.keys == ["sequence"]
def test_normalize_mixed_empty_key(): 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 = H5Store()
store._normalize({"sequence": [torch.tensor([1, 2, 3])], "loss_mask": []}) store._normalize({"sequence": [torch.tensor([1, 2, 3])], "loss_mask": []})
assert len(store) == 0 assert len(store) == 0
assert store.num_records == 0
assert store.token_count == 0 # min() over keys
assert set(store.keys) == {"sequence", "loss_mask"} 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.""" """GRPO dataset returns correct dtypes for per-record structured data."""
from astrai.dataset.dataset import GRPODataset from astrai.dataset.dataset import GRPODataset
test_dir = base_test_env["test_dir"]
G = 4 G = 4
dataset = GRPODataset() store = type(
dataset.storage = type(
"FakeStore", "FakeStore",
(), (),
{ {
"keys": ["prompts", "responses", "masks", "rewards"], "keys": ["prompts", "responses", "masks", "rewards"],
"num_records": 1, "num_records": 1,
"token_count": 0,
"_data": { "_data": {
"prompts": [torch.randint(0, 100, (10,), dtype=torch.int32)], "prompts": [torch.randint(0, 100, (10,), dtype=torch.int32)],
"responses": [ "responses": [
@@ -357,8 +359,10 @@ def test_grpo_dataset_dtype(base_test_env):
"rewards": [torch.rand(G, dtype=torch.float32)], "rewards": [torch.rand(G, dtype=torch.float32)],
}, },
"fetch_record": _fake_fetch_record, "fetch_record": _fake_fetch_record,
"__len__": lambda self: self.num_records,
}, },
)() )()
dataset = GRPODataset(store=store)
item = dataset[0] item = dataset[0]
assert item["prompts"].dtype == torch.long 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.""" """GRPO dataset loads record-structured data with per-response boundaries."""
from astrai.dataset.dataset import GRPODataset from astrai.dataset.dataset import GRPODataset
test_dir = base_test_env["test_dir"]
G = 3 G = 3
prompt_len = 8 prompt_len = 8
resp_lens = [5, 7, 4] resp_lens = [5, 7, 4]
dataset = GRPODataset() store = type(
dataset.storage = type(
"FakeStore", "FakeStore",
(), (),
{ {
"keys": ["prompts", "responses", "masks", "rewards"], "keys": ["prompts", "responses", "masks", "rewards"],
"num_records": 1, "num_records": 1,
"token_count": 0,
"_data": { "_data": {
"prompts": [torch.randint(0, 100, (prompt_len,))], "prompts": [torch.randint(0, 100, (prompt_len,))],
"responses": [[torch.randint(0, 100, (rl,)) for rl in resp_lens]], "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)], "rewards": [torch.tensor([0.9, 0.3, 0.7], dtype=torch.float32)],
}, },
"fetch_record": _fake_fetch_record, "fetch_record": _fake_fetch_record,
"__len__": lambda self: self.num_records,
}, },
)() )()
dataset = GRPODataset(store=store)
assert len(dataset) == 1 assert len(dataset) == 1
item = dataset[0] item = dataset[0]
@@ -458,7 +463,7 @@ def test_dataset_load_explicit_storage_type(base_test_env):
test_dir = base_test_env["test_dir"] test_dir = base_test_env["test_dir"]
dataset = _make_seq_dataset(test_dir, "explicit", storage_type="h5") dataset = _make_seq_dataset(test_dir, "explicit", storage_type="h5")
assert len(dataset) > 0 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): 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, 0] == 4
assert result["responses"][0, 0, 1] == 5 assert result["responses"][0, 0, 1] == 5
assert result["responses"][0, 0, 2] == 0 # padded 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 # Check response content: item 0, response 1 is [6,7,8,9] no padding
assert result["responses"][0, 1, 3] == 9 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): 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)] [torch.randint(0, 100, (np.random.randint(3, 8),)) for _ in range(G)]
for _ in range(n_records) for _ in range(n_records)
] ]
dataset = GRPODataset() store = type(
dataset.storage = type(
"FakeStore", "FakeStore",
(), (),
{ {
"keys": ["prompts", "responses", "masks", "rewards"], "keys": ["prompts", "responses", "masks", "rewards"],
"num_records": n_records, "num_records": n_records,
"token_count": 0,
"_data": { "_data": {
"prompts": [torch.randint(0, 100, (10,)) for _ in range(n_records)], "prompts": [torch.randint(0, 100, (10,)) for _ in range(n_records)],
"responses": dummy_responses, "responses": dummy_responses,
@@ -890,8 +895,10 @@ def test_grpo_multiple_records(base_test_env):
], ],
}, },
"fetch_record": _fake_fetch_record, "fetch_record": _fake_fetch_record,
"__len__": lambda self: self.num_records,
}, },
)() )()
dataset = GRPODataset(store=store)
assert len(dataset) == n_records assert len(dataset) == n_records
@@ -973,8 +980,8 @@ def test_dpo_jsonl_lazy_load(base_test_env):
) )
assert len(ds) == 2 assert len(ds) == 2
assert ds.storage.num_records == 2 assert ds.store.num_records == 2
assert ds.storage._processor is not None assert ds.store._processor is not None
item = ds[0] item = ds[0]
assert set(item.keys()) == {"chosen", "rejected", "chosen_mask", "rejected_mask"} 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): 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"] test_dir = base_test_env["test_dir"]
seq_length = 64 seq_length = 64
@@ -1048,8 +1061,9 @@ def test_h5_store_dual_mode(base_test_env):
store = H5Store() store = H5Store()
store.load(test_dir) store.load(test_dir)
assert len(store) == seq_length * 2 assert store.token_count == seq_length * 2
assert store.num_records == 2 assert store.num_records == 2
assert len(store) == 2 # no window configured → record count
rec0 = store.fetch_record(0, "chosen") rec0 = store.fetch_record(0, "chosen")
assert rec0.shape == (seq_length,) assert rec0.shape == (seq_length,)
@@ -1057,9 +1071,20 @@ def test_h5_store_dual_mode(base_test_env):
stream = store.fetch(0, 10, "chosen") stream = store.fetch(0, 10, "chosen")
assert stream.shape == (10,) 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): 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"] test_dir = base_test_env["test_dir"]
seq_length = 128 seq_length = 128
@@ -1069,8 +1094,9 @@ def test_mmap_store_stream_only_no_offsets(base_test_env):
store = StoreFactory.create("bin") store = StoreFactory.create("bin")
store.load(test_dir) store.load(test_dir)
assert len(store) == seq_length assert store.token_count == seq_length
assert store.num_records == 0 assert store.num_records == 0
assert len(store) == 0
chunk = store.fetch(0, 32, "sequence") chunk = store.fetch(0, 32, "sequence")
assert chunk.shape == (32,) assert chunk.shape == (32,)