refactor: move sample-id indexing from dataset to store
- 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:
+136
-241
@@ -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.
|
||||
|
||||
Implements common functionality for window-based data fetching.
|
||||
Uses a storage abstraction for format-agnostic data loading.
|
||||
"""
|
||||
|
||||
def __init__(self, window_size: int, stride: int):
|
||||
super().__init__()
|
||||
self.window_size = window_size
|
||||
self.stride = stride
|
||||
self.storage: Optional[Store] = None
|
||||
|
||||
@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:
|
||||
def validate_keys(store: Store, required: List[str]) -> None:
|
||||
"""Raise ``KeyError`` if *store* is missing any *required* key."""
|
||||
if not required:
|
||||
return
|
||||
actual_keys = set(self.storage.keys)
|
||||
missing = [k for k in self.required_keys if k not in actual_keys]
|
||||
actual = set(store.keys)
|
||||
missing = [k for k in required if k not in actual]
|
||||
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}"
|
||||
f"Store at {getattr(store, '_load_path', '?')} is missing required "
|
||||
f"keys {missing}; available keys are {sorted(actual)}."
|
||||
)
|
||||
|
||||
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.
|
||||
class BaseDataset(Dataset, ABC):
|
||||
"""Abstract base class for dataset types.
|
||||
|
||||
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.
|
||||
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).
|
||||
"""
|
||||
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)
|
||||
required_keys: List[str] = []
|
||||
|
||||
def __init__(self, store: Store):
|
||||
super().__init__()
|
||||
self.store: Store = store
|
||||
validate_keys(store, self.required_keys)
|
||||
|
||||
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],
|
||||
|
||||
+222
-76
@@ -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,15 +339,20 @@ class Streamable:
|
||||
end: int,
|
||||
keys: Union[str, List[str]],
|
||||
):
|
||||
if not self._data:
|
||||
return _stream_fetch(self, begin, end, keys)
|
||||
|
||||
|
||||
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 self._fetch_stream_key(keys, begin, end)
|
||||
return {k: self._fetch_stream_key(k, begin, end) for k in keys}
|
||||
return _fetch_stream_key(self, keys, begin, end)
|
||||
return {k: _fetch_stream_key(self, k, begin, end) for k in keys}
|
||||
|
||||
|
||||
def _fetch_stream_key(self, key: str, begin: int, end: int) -> Tensor:
|
||||
segments = self._data[key]
|
||||
@@ -243,33 +371,31 @@ class Streamable:
|
||||
|
||||
|
||||
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:
|
||||
return _record_fetch(self, index, keys)
|
||||
|
||||
|
||||
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 self._fetch_record_key(keys, index)
|
||||
return {k: self._fetch_record_key(k, index) for k in keys}
|
||||
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)
|
||||
@@ -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)
|
||||
|
||||
+71
-45
@@ -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,)
|
||||
|
||||
Reference in New Issue
Block a user