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:
+140
-245
@@ -1,21 +1,26 @@
|
||||
"""Dataset implementations with factory pattern for training.
|
||||
"""Dataset implementations for training.
|
||||
|
||||
Composition over inheritance — every dataset is a thin wrapper that
|
||||
binds a :class:`Store` to a particular train-type's key mapping. All
|
||||
sample-id → token/record indexing lives on the Store; datasets never
|
||||
know about window/stride math or segment layouts.
|
||||
|
||||
Class hierarchy:
|
||||
|
||||
BaseDataset (ABC) — load/validate, owns a Store
|
||||
├── SEQDataset — stream, next-token prediction (PT)
|
||||
├── SFTDataset — stream, loss-mask + position_ids
|
||||
└── RecordDataset — record access, optional processor
|
||||
├── DPODataset — chosen/rejected pairs
|
||||
└── GRPODataset — prompt + response group
|
||||
BaseDataset (ABC) — holds a Store, exposes __len__/keys,
|
||||
overrides __getitem__
|
||||
├── SEQDataset — next-token prediction (stream)
|
||||
├── SFTDataset — loss-mask + position_ids (stream)
|
||||
├── DPODataset — chosen/rejected pairs (record)
|
||||
└── GRPODataset — prompt + response group (record)
|
||||
|
||||
``RecordDataset`` holds an optional *processor* (pure
|
||||
``record -> Dict[str, Tensor]`` function). When the backing Store is
|
||||
a lazy JsonlStore, the processor tokenises on the fly; otherwise it
|
||||
is ignored and ``fetch_record`` reads pre-tokenised tensors.
|
||||
``DatasetFactory.load(train_type, load_path, window_size, stride, …)``
|
||||
builds the Store (auto-detecting format) before constructing the
|
||||
matching dataset. Passing ``store=`` skips Store construction.
|
||||
|
||||
``__len__`` returns the sample count (stream: windows, record:
|
||||
records) so DataLoader and progress bars work uniformly.
|
||||
When a record dataset (DPO) reads from raw JSONL, a *processor*
|
||||
function (pure ``record -> Dict[str, Tensor]``) is forwarded to
|
||||
:class:`JsonlStore` so tokenisation happens on the fly.
|
||||
"""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
@@ -218,184 +223,57 @@ def grpo_collate_fn(batch: List[Dict[str, Tensor]]) -> Dict[str, Tensor]:
|
||||
}
|
||||
|
||||
|
||||
class BaseDataset(Dataset, ABC):
|
||||
"""Abstract base class for all dataset types.
|
||||
def validate_keys(store: Store, required: List[str]) -> None:
|
||||
"""Raise ``KeyError`` if *store* is missing any *required* key."""
|
||||
if not required:
|
||||
return
|
||||
actual = set(store.keys)
|
||||
missing = [k for k in required if k not in actual]
|
||||
if missing:
|
||||
raise KeyError(
|
||||
f"Store at {getattr(store, '_load_path', '?')} is missing required "
|
||||
f"keys {missing}; available keys are {sorted(actual)}."
|
||||
)
|
||||
|
||||
Implements common functionality for window-based data fetching.
|
||||
Uses a storage abstraction for format-agnostic data loading.
|
||||
|
||||
class BaseDataset(Dataset, ABC):
|
||||
"""Abstract base class for dataset types.
|
||||
|
||||
Holds a :class:`Store`. All sample-id indexing is delegated to the
|
||||
store — this class exposes ``__len__`` as ``len(store)`` and the
|
||||
``keys`` property as ``store.keys``. Subclasses implement
|
||||
``__getitem__`` with the train-type-specific key mapping and any
|
||||
training-only index arithmetic (e.g. the next-token ``+1`` shift).
|
||||
"""
|
||||
|
||||
def __init__(self, window_size: int, stride: int):
|
||||
required_keys: List[str] = []
|
||||
|
||||
def __init__(self, store: Store):
|
||||
super().__init__()
|
||||
self.window_size = window_size
|
||||
self.stride = stride
|
||||
self.storage: Optional[Store] = None
|
||||
self.store: Store = store
|
||||
validate_keys(store, self.required_keys)
|
||||
|
||||
@property
|
||||
def required_keys(self) -> List[str]:
|
||||
"""Return required storage keys for this dataset type.
|
||||
|
||||
Subclasses should override to specify expected keys.
|
||||
"""
|
||||
return []
|
||||
|
||||
def _validate_keys(self):
|
||||
if not self.required_keys:
|
||||
return
|
||||
actual_keys = set(self.storage.keys)
|
||||
missing = [k for k in self.required_keys if k not in actual_keys]
|
||||
if missing:
|
||||
raise KeyError(
|
||||
f"Dataset {type(self).__name__} requires keys {self.required_keys}, "
|
||||
f"but storage at {self._load_path} only has {sorted(actual_keys)}. "
|
||||
f"Missing: {missing}"
|
||||
)
|
||||
|
||||
def load(self, load_path: str, storage_type: Optional[str] = None, **kwargs):
|
||||
"""Load dataset from the given path.
|
||||
|
||||
Auto-detects the storage format if not specified.
|
||||
|
||||
Args:
|
||||
load_path: Path to the data directory or file
|
||||
storage_type: Force a specific storage type ("h5", "bin", "jsonl"),
|
||||
or None for auto-detection
|
||||
**kwargs: Extra arguments forwarded to the store constructor and
|
||||
to ``store.load()``.
|
||||
|
||||
Raises:
|
||||
KeyError: If the loaded storage is missing required keys.
|
||||
"""
|
||||
if storage_type is None:
|
||||
storage_type = detect_format(load_path)
|
||||
self.storage = StoreFactory.create(storage_type, **kwargs)
|
||||
self._load_path = load_path
|
||||
self.storage.load(load_path, **kwargs)
|
||||
self._validate_keys()
|
||||
|
||||
@property
|
||||
def count(self) -> int:
|
||||
"""Return the total number of raw elements (tokens) in the dataset."""
|
||||
if self.storage is None:
|
||||
return 0
|
||||
return len(self.storage)
|
||||
def __len__(self) -> int:
|
||||
return len(self.store)
|
||||
|
||||
@property
|
||||
def keys(self) -> List[str]:
|
||||
"""Return the available data keys."""
|
||||
if self.storage is None:
|
||||
return []
|
||||
return self.storage.keys
|
||||
return self.store.keys
|
||||
|
||||
def get_index(self, index: int) -> tuple:
|
||||
"""Calculate begin and end indices for a sample.
|
||||
|
||||
Args:
|
||||
index: Sample index
|
||||
|
||||
Returns:
|
||||
Tuple of (begin_idx, end_idx)
|
||||
"""
|
||||
if self.storage is None:
|
||||
raise RuntimeError("Dataset not loaded, call load() first")
|
||||
total = len(self.storage)
|
||||
if total <= self.window_size:
|
||||
raise ValueError(
|
||||
f"Data too short: {total} tokens <= window_size {self.window_size}"
|
||||
)
|
||||
|
||||
begin_idx = min(index * self.stride, total - 1 - self.window_size)
|
||||
end_idx = min(begin_idx + self.window_size, total - 1)
|
||||
|
||||
return begin_idx, end_idx
|
||||
@property
|
||||
def token_count(self) -> int:
|
||||
return self.store.token_count
|
||||
|
||||
@abstractmethod
|
||||
def __getitem__(self, index: int) -> Dict[str, Tensor]:
|
||||
"""Get a single sample by index.
|
||||
|
||||
Must be implemented by subclasses.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def __len__(self) -> int:
|
||||
if self.storage is None:
|
||||
return 0
|
||||
total = len(self.storage)
|
||||
if total <= self.window_size:
|
||||
return 0
|
||||
return (total - 1 - self.window_size) // self.stride + 1
|
||||
|
||||
|
||||
class RecordDataset(BaseDataset):
|
||||
"""Base class for record-structured datasets (DPO/GRPO).
|
||||
|
||||
Each sample is an independent record — no windowing, stride, or
|
||||
cross-record concatenation. ``__len__`` returns the record count
|
||||
so progress bars advance per-record.
|
||||
|
||||
A *processor* (pure ``record -> Dict[str, Tensor]`` function) may be
|
||||
supplied for lazy on-the-fly tokenisation of raw JSONL. The
|
||||
processor is forwarded to ``JsonlStore`` and applied per access;
|
||||
pre-tokenised backends (H5/bin) ignore it.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
window_size: int = 0,
|
||||
stride: int = 0,
|
||||
processor: Optional[Callable[[dict], Dict[str, Tensor]]] = None,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__(window_size=window_size, stride=stride or window_size)
|
||||
self.processor = processor
|
||||
|
||||
def load(self, load_path: str, storage_type: Optional[str] = None, **kwargs):
|
||||
"""Load data from *load_path*.
|
||||
|
||||
Args:
|
||||
load_path: Path to data file or directory.
|
||||
storage_type: Force backend ("h5"/"bin"/"jsonl") or None for
|
||||
auto-detection.
|
||||
**kwargs: Forwarded to ``store.load()``. When the backend is
|
||||
JSONL and a processor was set, it is passed as
|
||||
``processor=`` for lazy tokenisation.
|
||||
"""
|
||||
if storage_type is None:
|
||||
storage_type = detect_format(load_path)
|
||||
self.storage = StoreFactory.create(storage_type, **kwargs)
|
||||
self._load_path = load_path
|
||||
|
||||
if self.processor is not None:
|
||||
self.storage.load(load_path, processor=self.processor, **kwargs)
|
||||
else:
|
||||
self.storage.load(load_path, **kwargs)
|
||||
self._validate_keys()
|
||||
|
||||
def __len__(self) -> int:
|
||||
if self.storage is None:
|
||||
return 0
|
||||
return self.storage.num_records
|
||||
|
||||
@property
|
||||
def count(self) -> int:
|
||||
if self.storage is None:
|
||||
return 0
|
||||
return self.storage.num_records
|
||||
|
||||
|
||||
class DatasetFactory(BaseFactory["BaseDataset"]):
|
||||
"""Factory class for creating dataset instances.
|
||||
"""Factory for creating dataset instances by train-type.
|
||||
|
||||
Supports decorator-based registration for extensible dataset types.
|
||||
All default dataset types (seq, sft, dpo, grpo) are registered automatically
|
||||
when their classes are defined with the decorator.
|
||||
|
||||
Example usage:
|
||||
@DatasetFactory.register("custom")
|
||||
class CustomDataset(BaseDataset):
|
||||
...
|
||||
|
||||
dataset = DatasetFactory.create("custom", window_size, stride)
|
||||
Use :meth:`DatasetFactory.register("custom")` to register new
|
||||
dataset classes; they must inherit from :class:`BaseDataset`.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
@@ -417,32 +295,34 @@ class DatasetFactory(BaseFactory["BaseDataset"]):
|
||||
|
||||
- **store given**: bind it directly — the caller fully controls
|
||||
Store construction and processor setup. *load_path*,
|
||||
*storage_type*, *tokenizer_path* are ignored.
|
||||
*storage_type*, *tokenizer_path*, *window_size*, *stride* are
|
||||
ignored.
|
||||
- **store is None**: build a Store from *load_path*, auto-detecting
|
||||
format and constructing a processor when *tokenizer_path* is
|
||||
given for a record dataset on JSONL.
|
||||
|
||||
Args:
|
||||
train_type: Type of training dataset
|
||||
load_path: Path to the data file (ignored if *store* given)
|
||||
window_size: Window size for data sampling
|
||||
stride: Stride between consecutive samples (default: same as window_size)
|
||||
storage_type: Storage type ("h5", "bin", "jsonl") or None for auto-detection
|
||||
tokenizer_path: Path to tokenizer for lazy JSONL tokenisation
|
||||
max_len: Max sequence length for the processor
|
||||
store: Pre-built, already-loaded Store instance
|
||||
**kwargs: Extra arguments forwarded to ``dataset.load()``
|
||||
train_type: Registered dataset name ("seq", "sft", "dpo",
|
||||
"grpo", …).
|
||||
load_path: Path to the data file or directory (ignored if
|
||||
*store* is given).
|
||||
window_size: Stream window length — only meaningful for
|
||||
stream datasets (SEQ/SFT). Record datasets ignore it.
|
||||
stride: Stride between consecutive stream samples
|
||||
(default: same as *window_size*).
|
||||
storage_type: Storage backend ("h5", "bin", "jsonl") or
|
||||
None for auto-detection.
|
||||
tokenizer_path: Path to tokenizer for lazy JSONL
|
||||
tokenisation (record datasets only).
|
||||
max_len: Max sequence length forwarded to processors.
|
||||
store: Pre-built, already-loaded Store instance.
|
||||
**kwargs: Extra arguments forwarded to ``store.load()``.
|
||||
|
||||
Returns:
|
||||
Loaded dataset instance
|
||||
Loaded dataset instance.
|
||||
"""
|
||||
if stride is None:
|
||||
stride = window_size
|
||||
|
||||
if store is not None:
|
||||
dataset = cls.create(train_type, window_size, stride)
|
||||
dataset.storage = store
|
||||
return dataset
|
||||
return cls.create(train_type, store=store)
|
||||
|
||||
if load_path is None:
|
||||
raise ValueError("Either load_path or store must be provided")
|
||||
@@ -450,14 +330,37 @@ class DatasetFactory(BaseFactory["BaseDataset"]):
|
||||
if storage_type is None:
|
||||
storage_type = detect_format(load_path)
|
||||
|
||||
if stride is None:
|
||||
stride = window_size
|
||||
|
||||
processor = cls._maybe_build_processor(
|
||||
train_type, storage_type, tokenizer_path, max_len
|
||||
)
|
||||
|
||||
dataset = cls.create(train_type, window_size, stride, processor=processor)
|
||||
dataset.load(load_path, storage_type=storage_type, **kwargs)
|
||||
store_window = cls._store_window_for(train_type, window_size)
|
||||
store = StoreFactory.create(
|
||||
storage_type,
|
||||
window_size=store_window,
|
||||
stride=stride if stride else store_window,
|
||||
)
|
||||
if processor is not None:
|
||||
store.load(load_path, processor=processor, **kwargs)
|
||||
else:
|
||||
store.load(load_path, **kwargs)
|
||||
|
||||
return dataset
|
||||
return cls.create(train_type, store=store)
|
||||
|
||||
@staticmethod
|
||||
def _store_window_for(train_type: str, window_size: int) -> int:
|
||||
"""Stream datasets consume ``window_size``; record datasets ignore it.
|
||||
|
||||
Record datasets (dpo/grpo) treat each record as an independent
|
||||
training unit and never window, so the store is built with
|
||||
``window_size=0`` and ``len(store)`` returns the record count.
|
||||
"""
|
||||
if train_type in ("seq", "sft"):
|
||||
return window_size
|
||||
return 0
|
||||
|
||||
@staticmethod
|
||||
def _maybe_build_processor(
|
||||
@@ -482,43 +385,41 @@ class DatasetFactory(BaseFactory["BaseDataset"]):
|
||||
|
||||
@DatasetFactory.register("seq")
|
||||
class SEQDataset(BaseDataset):
|
||||
"""Dataset for sequential next-token prediction training."""
|
||||
"""Dataset for sequential next-token prediction training.
|
||||
|
||||
@property
|
||||
def required_keys(self) -> List[str]:
|
||||
return ["sequence"]
|
||||
Stream mode: ``store.fetch(begin, end, "sequence")`` returns the
|
||||
input window; the +1 shifted call returns the next-token target.
|
||||
"""
|
||||
|
||||
def _fetch_data(self, begin_idx: int, end_idx: int) -> Tensor:
|
||||
return self.storage.fetch(begin_idx, end_idx, "sequence")
|
||||
required_keys = ["sequence"]
|
||||
|
||||
def __getitem__(self, index):
|
||||
begin_idx, end_idx = self.get_index(index)
|
||||
|
||||
x = self._fetch_data(begin_idx, end_idx).to(dtype=torch.long)
|
||||
y = self._fetch_data(begin_idx + 1, end_idx + 1).to(dtype=torch.long)
|
||||
|
||||
return {"input_ids": x, "target_ids": y}
|
||||
def __getitem__(self, index: int):
|
||||
begin, end = self.store.sample_window(index)
|
||||
x = self.store.fetch(begin, end, "sequence")
|
||||
y = self.store.fetch(begin + 1, end + 1, "sequence")
|
||||
return {
|
||||
"input_ids": x.to(dtype=torch.long),
|
||||
"target_ids": y.to(dtype=torch.long),
|
||||
}
|
||||
|
||||
|
||||
@DatasetFactory.register("sft")
|
||||
class SFTDataset(BaseDataset):
|
||||
"""Dataset for supervised fine-tuning with loss masking."""
|
||||
"""Dataset for supervised fine-tuning with loss masking.
|
||||
|
||||
@property
|
||||
def required_keys(self) -> List[str]:
|
||||
return ["sequence", "loss_mask", "position_ids"]
|
||||
Stream mode: ``sequence``/``loss_mask``/``position_ids`` are sliced
|
||||
to the window. ``loss_mask`` and ``target_ids`` use the +1 shifted
|
||||
slice so they align with the predicted positions.
|
||||
"""
|
||||
|
||||
def _fetch_data(self, begin_idx: int, end_idx: int, key: str) -> Tensor:
|
||||
return self.storage.fetch(begin_idx, end_idx, key)
|
||||
|
||||
def __getitem__(self, index):
|
||||
begin_idx, end_idx = self.get_index(index)
|
||||
|
||||
x = self._fetch_data(begin_idx, end_idx, "sequence")
|
||||
y = self._fetch_data(begin_idx + 1, end_idx + 1, "sequence")
|
||||
position_ids = self._fetch_data(begin_idx, end_idx, "position_ids")
|
||||
loss_mask = self._fetch_data(begin_idx + 1, end_idx + 1, "loss_mask")
|
||||
required_keys = ["sequence", "loss_mask", "position_ids"]
|
||||
|
||||
def __getitem__(self, index: int):
|
||||
begin, end = self.store.sample_window(index)
|
||||
x = self.store.fetch(begin, end, "sequence")
|
||||
y = self.store.fetch(begin + 1, end + 1, "sequence")
|
||||
position_ids = self.store.fetch(begin, end, "position_ids")
|
||||
loss_mask = self.store.fetch(begin + 1, end + 1, "loss_mask")
|
||||
return {
|
||||
"input_ids": x.to(dtype=torch.long),
|
||||
"target_ids": y.to(dtype=torch.long),
|
||||
@@ -528,7 +429,7 @@ class SFTDataset(BaseDataset):
|
||||
|
||||
|
||||
@DatasetFactory.register("dpo")
|
||||
class DPODataset(RecordDataset):
|
||||
class DPODataset(BaseDataset):
|
||||
"""Record-structured dataset for Direct Preference Optimization.
|
||||
|
||||
Each sample is one preference pair (chosen + rejected) and is an
|
||||
@@ -536,39 +437,35 @@ class DPODataset(RecordDataset):
|
||||
concatenation. This keeps each sequence self-contained so attention
|
||||
never leaks across preference pairs.
|
||||
|
||||
Two loading paths (handled by :class:`RecordDataset`):
|
||||
Two loading paths (handled by :class:`DatasetFactory`):
|
||||
|
||||
- **Pre-tokenized** (H5/bin): ``load(path)`` reads per-record tensors,
|
||||
``__getitem__`` returns them directly.
|
||||
- **Raw JSONL** (``tokenizer_path=...``): builds a lazy processor via
|
||||
:func:`dpo_processor` that tokenises on the fly — no packing, no
|
||||
``position_ids``.
|
||||
- **Pre-tokenized** (H5/bin): ``store.load(path)`` reads per-record
|
||||
tensors; ``__getitem__`` returns them directly.
|
||||
- **Raw JSONL** (``tokenizer_path=...``): builds a lazy processor
|
||||
via :func:`dpo_processor` that tokenises on the fly — no packing,
|
||||
no ``position_ids``.
|
||||
"""
|
||||
|
||||
@property
|
||||
def required_keys(self) -> List[str]:
|
||||
return ["chosen", "rejected", "chosen_mask", "rejected_mask"]
|
||||
required_keys = ["chosen", "rejected", "chosen_mask", "rejected_mask"]
|
||||
|
||||
def make_processor(self, tokenizer, max_len: int):
|
||||
return partial(dpo_processor, tokenizer=tokenizer, max_len=max_len)
|
||||
|
||||
def __getitem__(self, index: int) -> Dict[str, Tensor]:
|
||||
return {
|
||||
"chosen": self.storage.fetch_record(index, "chosen").to(dtype=torch.long),
|
||||
"rejected": self.storage.fetch_record(index, "rejected").to(
|
||||
dtype=torch.long
|
||||
),
|
||||
"chosen_mask": self.storage.fetch_record(index, "chosen_mask").to(
|
||||
"chosen": self.store.fetch_record(index, "chosen").to(dtype=torch.long),
|
||||
"rejected": self.store.fetch_record(index, "rejected").to(dtype=torch.long),
|
||||
"chosen_mask": self.store.fetch_record(index, "chosen_mask").to(
|
||||
dtype=torch.bool
|
||||
),
|
||||
"rejected_mask": self.storage.fetch_record(index, "rejected_mask").to(
|
||||
"rejected_mask": self.store.fetch_record(index, "rejected_mask").to(
|
||||
dtype=torch.bool
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
@DatasetFactory.register("grpo")
|
||||
class GRPODataset(RecordDataset):
|
||||
class GRPODataset(BaseDataset):
|
||||
"""Dataset for offline Group Relative Policy Optimization.
|
||||
|
||||
Each sample is one prompt with its group of responses and scalar
|
||||
@@ -582,15 +479,13 @@ class GRPODataset(RecordDataset):
|
||||
- ``rewards``: List[Tensor] — one 1-D float tensor (len G) per record
|
||||
"""
|
||||
|
||||
@property
|
||||
def required_keys(self) -> List[str]:
|
||||
return ["prompts", "responses", "masks", "rewards"]
|
||||
required_keys = ["prompts", "responses", "masks", "rewards"]
|
||||
|
||||
def __getitem__(self, index: int) -> Dict[str, Tensor]:
|
||||
prompts = self.storage.fetch_record(index, "prompts")
|
||||
responses = self.storage.fetch_record(index, "responses")
|
||||
masks = self.storage.fetch_record(index, "masks")
|
||||
rewards = self.storage.fetch_record(index, "rewards")
|
||||
prompts = self.store.fetch_record(index, "prompts")
|
||||
responses = self.store.fetch_record(index, "responses")
|
||||
masks = self.store.fetch_record(index, "masks")
|
||||
rewards = self.store.fetch_record(index, "rewards")
|
||||
return {
|
||||
"prompts": prompts.to(dtype=torch.long),
|
||||
"responses": [r.to(dtype=torch.long) for r in responses],
|
||||
|
||||
+257
-111
@@ -1,11 +1,14 @@
|
||||
"""Storage backends for different data formats.
|
||||
|
||||
Architecture (mixin composition, no diamond inheritance):
|
||||
Architecture (composition over inheritance):
|
||||
|
||||
Store (ABC) — shared _data/_cum/_offsets bookkeeping
|
||||
+ _normalize() for registering segments
|
||||
Streamable (mixin) — fetch(begin, end, key) for stream access
|
||||
Recordable (mixin) — fetch_record(i, key) for record access
|
||||
Store (ABC) — owns _data/_cum/_offsets bookkeeping
|
||||
+ window_size/stride for sample-id
|
||||
indexing. __getitem__/__len__ produce
|
||||
the smallest iterable unit so Dataset
|
||||
classes are pure delegators.
|
||||
Streamable (mixin) — raw token slice fetch(begin, end, keys)
|
||||
Recordable (mixin) — raw record slice fetch_record(idx, keys)
|
||||
|
||||
H5Store(Store, Streamable, Recordable)
|
||||
MmapStore(Store, Streamable, Recordable)
|
||||
@@ -13,16 +16,24 @@ Architecture (mixin composition, no diamond inheritance):
|
||||
|
||||
Each mixin is a stateless trait that relies on ``self._data`` etc.
|
||||
provided by :class:`Store`. Concrete stores mix in whichever access
|
||||
modes they support — ``Store`` is the sole base class, so there is no
|
||||
diamond inheritance or MRO ambiguity.
|
||||
primitives they support — ``Store`` is the sole base class, so there is
|
||||
no diamond inheritance or MRO ambiguity.
|
||||
|
||||
Access-mode semantics:
|
||||
Sample-id indexing lives on :class:`Store`, not on the dataset:
|
||||
|
||||
- **Stream** (SEQ/SFT): ``fetch(begin, end, key)`` slices across
|
||||
concatenated segments. ``len(store)`` returns the total token count.
|
||||
- **Record** (DPO/GRPO): ``fetch_record(i, key)`` returns the *i*-th
|
||||
record without cross-record concatenation. ``num_records`` returns
|
||||
the record count.
|
||||
- **Stream mode** (``window_size > 0``): ``len(store)`` returns the number
|
||||
of ``(window_size, stride)`` windows that fit in the token river;
|
||||
``store[i]`` returns the *i*-th window as a dict of per-key tensors;
|
||||
``store.sample_window(i)`` exposes the underlying ``(begin, end)``
|
||||
token slice for callers (e.g. next-token trainers) that need a +1
|
||||
shifted companion window.
|
||||
- **Record mode** (``num_records > 0``): ``len(store)`` returns the
|
||||
record count; ``store[i]`` returns the *i*-th record dict.
|
||||
|
||||
Raw token/record access via :meth:`fetch` / :meth:`fetch_record`
|
||||
remains available for low-level callers that want explicit index
|
||||
control. ``store.token_count`` is the total stream token count (what
|
||||
``len(store)`` used to mean in the legacy stream-only API).
|
||||
|
||||
``segments_are_records`` (class attribute on each Store subclass)
|
||||
tells ``_normalize`` whether segments are inherently per-record (H5/
|
||||
@@ -40,7 +51,7 @@ import json
|
||||
import logging
|
||||
from abc import ABC, abstractmethod
|
||||
from pathlib import Path
|
||||
from typing import Callable, Dict, List, Optional, Union
|
||||
from typing import Callable, Dict, List, Optional, Tuple, Union
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
@@ -102,20 +113,46 @@ def detect_format(load_path: str) -> str:
|
||||
class Store(ABC):
|
||||
"""Common base for all storage backends.
|
||||
|
||||
Owns the shared ``_data`` / ``_cum`` / ``_offsets`` bookkeeping and
|
||||
the ``_normalize`` entry point used by tensor-backed subclasses.
|
||||
Does **not** expose an access API — that is the job of
|
||||
:class:`StreamStore` and :class:`RecordStore`.
|
||||
A Store owns both its data layout AND its sample-id → token/record
|
||||
index translation. Datasets are thin wrappers that bind a Store
|
||||
to a particular train-type's key mapping; they never know about
|
||||
window/stride math.
|
||||
|
||||
Two iteration modes:
|
||||
|
||||
- **Stream** (``window_size > 0``): data is treated as one long
|
||||
token river. ``len(store)`` returns the number of windows;
|
||||
``store[i]`` slices every stream-compatible key to window ``i``;
|
||||
``store.sample_window(i)`` returns the ``(begin, end)`` token
|
||||
slice for callers needing a +1 shifted companion window.
|
||||
- **Record** (``num_records > 0``): data is per-record.
|
||||
``len(store)`` returns ``num_records``; ``store[i]`` returns
|
||||
the *i*-th record as a dict.
|
||||
|
||||
Raw token slicing is still available via :meth:`fetch` (mixed in
|
||||
by :class:`Streamable`) when a store has stream support configured.
|
||||
Raw record slicing via :meth:`fetch_record` (mixed in by
|
||||
:class:`Recordable`) when a store has record support.
|
||||
|
||||
``token_count`` exposes the raw total stream length — this is what
|
||||
``len(store)`` returned in the legacy stream-only API and what
|
||||
stream-bound ``fetch`` uses for its bounds check.
|
||||
"""
|
||||
|
||||
segments_are_records: bool = False
|
||||
|
||||
def __init__(self):
|
||||
def __init__(
|
||||
self,
|
||||
window_size: int = 0,
|
||||
stride: Optional[int] = None,
|
||||
):
|
||||
self._data: Dict[str, List[Tensor]] = {}
|
||||
self._cum: Dict[str, List[int]] = {}
|
||||
self._offsets: Dict[str, List[int]] = {}
|
||||
self._length: int = 0
|
||||
self._num_records: int = 0
|
||||
self._window_size: int = int(window_size)
|
||||
self._stride: int = int(stride) if stride is not None else int(window_size)
|
||||
|
||||
@abstractmethod
|
||||
def load(self, path: str, **kwargs) -> None:
|
||||
@@ -125,14 +162,98 @@ class Store(ABC):
|
||||
def keys(self) -> List[str]:
|
||||
return list(self._data.keys())
|
||||
|
||||
def __len__(self) -> int:
|
||||
"""Default: token count (stream semantics).
|
||||
@property
|
||||
def window_size(self) -> int:
|
||||
return self._window_size
|
||||
|
||||
Subclasses that are record-only (e.g. lazy JsonlStore) override
|
||||
to return ``self._num_records``.
|
||||
@property
|
||||
def stride(self) -> int:
|
||||
return self._stride
|
||||
|
||||
@property
|
||||
def token_count(self) -> int:
|
||||
"""Total tokens across all stream segments.
|
||||
|
||||
Useful for the bounds-checked raw :meth:`fetch` and as the
|
||||
legacy ``len(store)`` value.
|
||||
"""
|
||||
return self._length
|
||||
|
||||
@property
|
||||
def num_records(self) -> int:
|
||||
"""Number of records available via :meth:`fetch_record`.
|
||||
|
||||
Non-zero only when the backing layout provides per-record
|
||||
indexing (H5/JSONL segments or bin ``_offsets``).
|
||||
"""
|
||||
return self._num_records
|
||||
|
||||
@property
|
||||
def num_samples(self) -> int:
|
||||
"""Number of items produced by ``__getitem__``.
|
||||
|
||||
Stream-mode wins when ``window_size > 0`` and there are tokens
|
||||
to slice; otherwise falls back to ``num_records``.
|
||||
"""
|
||||
if self._window_size > 0 and self._length > 0:
|
||||
total = self._length
|
||||
w = self._window_size
|
||||
if total <= w:
|
||||
return 0
|
||||
return (total - 1 - w) // self._stride + 1
|
||||
return self._num_records
|
||||
|
||||
def __len__(self) -> int:
|
||||
return self.num_samples
|
||||
|
||||
def __getitem__(self, index: int) -> Dict[str, Tensor]:
|
||||
if index < 0:
|
||||
index += self.num_samples
|
||||
if not 0 <= index < self.num_samples:
|
||||
raise IndexError(
|
||||
f"Store index out of range: {index}, num_samples={self.num_samples}"
|
||||
)
|
||||
if self._window_size > 0 and self._length > 0:
|
||||
begin, end = self.sample_window(index)
|
||||
keys = self._stream_keys()
|
||||
return {k: self.fetch(begin, end, k) for k in keys}
|
||||
return self.fetch_record(index, self._record_keys())
|
||||
|
||||
def sample_window(self, index: int) -> Tuple[int, int]:
|
||||
"""Return ``(begin, end)`` token positions for stream sample *index*.
|
||||
|
||||
The clipped tail keeps the last reachable window inside the
|
||||
token river instead of overshooting. Caller is responsible
|
||||
for staying within :attr:`num_samples`: an out-of-range index
|
||||
raises ``IndexError``.
|
||||
"""
|
||||
if self._window_size <= 0:
|
||||
raise RuntimeError("sample_window() requires window_size > 0 (stream mode)")
|
||||
if self._window_size <= 0 or self._length <= self._window_size:
|
||||
raise IndexError(
|
||||
f"Data too short for window: token_count={self._length}, "
|
||||
f"window_size={self._window_size}"
|
||||
)
|
||||
if not 0 <= index < self.num_samples:
|
||||
raise IndexError(
|
||||
f"Sample index out of range: {index}, num_samples={self.num_samples}"
|
||||
)
|
||||
total = self._length
|
||||
begin = min(index * self._stride, total - 1 - self._window_size)
|
||||
end = min(begin + self._window_size, total - 1)
|
||||
return begin, end
|
||||
|
||||
def _stream_keys(self) -> List[str]:
|
||||
out: List[str] = []
|
||||
for k, tensors in self._data.items():
|
||||
if tensors and isinstance(tensors[0], list):
|
||||
continue
|
||||
out.append(k)
|
||||
return out
|
||||
|
||||
def _record_keys(self) -> List[str]:
|
||||
return list(self._data.keys())
|
||||
|
||||
def _normalize(
|
||||
self,
|
||||
raw: Dict[str, list],
|
||||
@@ -141,14 +262,14 @@ class Store(ABC):
|
||||
"""Register segments and pre-compute indices for both access modes.
|
||||
|
||||
Stream mode: ``_cum[key]`` accumulates per-segment lengths so
|
||||
``StreamStore._fetch_key`` can bisect across segments without
|
||||
concatenation.
|
||||
``Streamable._fetch_stream_key`` can bisect across segments
|
||||
without concatenation.
|
||||
|
||||
Record mode: if *offsets* is provided (bin layout),
|
||||
``_offsets[key]`` stores cumulative per-record offsets into the
|
||||
single concatenated segment. Otherwise, when
|
||||
``segments_are_records`` is True (H5/JSONL), ``_data[key]`` is a
|
||||
per-record list and ``fetch_record`` indexes it directly.
|
||||
``segments_are_records`` is True (H5/JSONL), ``_data[key]`` is
|
||||
a per-record list and ``fetch_record`` indexes it directly.
|
||||
|
||||
Nested keys (GRPO ``responses``/``masks`` as
|
||||
``List[List[Tensor]]``) are stored as-is and excluded from both
|
||||
@@ -194,7 +315,7 @@ class Store(ABC):
|
||||
elif self.segments_are_records:
|
||||
per_record_counts = []
|
||||
for key, tensors in self._data.items():
|
||||
if not tensors or isinstance(tensors[0], list):
|
||||
if tensors and isinstance(tensors[0], list):
|
||||
continue
|
||||
per_record_counts.append(len(tensors))
|
||||
self._num_records = min(per_record_counts) if per_record_counts else 0
|
||||
@@ -203,11 +324,13 @@ class Store(ABC):
|
||||
|
||||
|
||||
class Streamable:
|
||||
"""Mixin: stream access ``fetch(begin, end, key)``.
|
||||
"""Mixin granting raw token-stream access via :meth:`fetch`.
|
||||
|
||||
No base class — relies on ``self._data``, ``self._cum``,
|
||||
``self._length`` provided by :class:`Store`. Used by SEQ/SFT
|
||||
where data is a long token stream.
|
||||
Stateless trait relying on ``self._data``, ``self._cum``,
|
||||
``self._length`` maintained by :class:`Store`. Stream mode is
|
||||
active when the owning store has ``window_size > 0``; for stores
|
||||
that can also serve record access (H5/JSONL/bin+offsets), the
|
||||
``fetch_record`` API from :class:`Recordable` is used instead.
|
||||
"""
|
||||
|
||||
def fetch(
|
||||
@@ -216,72 +339,75 @@ class Streamable:
|
||||
end: int,
|
||||
keys: Union[str, List[str]],
|
||||
):
|
||||
if not self._data:
|
||||
raise RuntimeError("Store not loaded")
|
||||
if not (0 <= begin < self._length and 0 <= end <= self._length):
|
||||
raise ValueError(
|
||||
f"Index out of bounds: begin={begin}, end={end}, length={self._length}"
|
||||
)
|
||||
if isinstance(keys, str):
|
||||
return self._fetch_stream_key(keys, begin, end)
|
||||
return {k: self._fetch_stream_key(k, begin, end) for k in keys}
|
||||
return _stream_fetch(self, begin, end, keys)
|
||||
|
||||
def _fetch_stream_key(self, key: str, begin: int, end: int) -> Tensor:
|
||||
segments = self._data[key]
|
||||
cum = self._cum[key]
|
||||
seg_start = bisect.bisect_right(cum, begin)
|
||||
seg_end = bisect.bisect_left(cum, end)
|
||||
|
||||
results = []
|
||||
for i in range(seg_start, seg_end + 1):
|
||||
prev = cum[i - 1] if i > 0 else 0
|
||||
s = max(begin - prev, 0)
|
||||
e = min(end - prev, segments[i].shape[0])
|
||||
results.append(segments[i][s:e])
|
||||
def _stream_fetch(self, begin: int, end: int, keys: Union[str, List[str]]):
|
||||
if not getattr(self, "_data", None):
|
||||
raise RuntimeError("Store not loaded")
|
||||
if not (0 <= begin < self._length and 0 <= end <= self._length):
|
||||
raise ValueError(
|
||||
f"Index out of bounds: begin={begin}, end={end}, length={self._length}"
|
||||
)
|
||||
if isinstance(keys, str):
|
||||
return _fetch_stream_key(self, keys, begin, end)
|
||||
return {k: _fetch_stream_key(self, k, begin, end) for k in keys}
|
||||
|
||||
return results[0] if len(results) == 1 else torch.cat(results, dim=0)
|
||||
|
||||
def _fetch_stream_key(self, key: str, begin: int, end: int) -> Tensor:
|
||||
segments = self._data[key]
|
||||
cum = self._cum[key]
|
||||
seg_start = bisect.bisect_right(cum, begin)
|
||||
seg_end = bisect.bisect_left(cum, end)
|
||||
|
||||
results = []
|
||||
for i in range(seg_start, seg_end + 1):
|
||||
prev = cum[i - 1] if i > 0 else 0
|
||||
s = max(begin - prev, 0)
|
||||
e = min(end - prev, segments[i].shape[0])
|
||||
results.append(segments[i][s:e])
|
||||
|
||||
return results[0] if len(results) == 1 else torch.cat(results, dim=0)
|
||||
|
||||
|
||||
class Recordable:
|
||||
"""Mixin: record access ``fetch_record(i, key)``.
|
||||
"""Mixin granting raw record access via :meth:`fetch_record`.
|
||||
|
||||
No base class — relies on ``self._data``, ``self._offsets``,
|
||||
``self._num_records`` provided by :class:`Store`. Used by
|
||||
DPO/GRPO where each record is an independent training unit.
|
||||
Stateless trait relying on ``self._data``, ``self._offsets``,
|
||||
``self._num_records`` maintained by :class:`Store`.
|
||||
"""
|
||||
|
||||
segments_are_records = True
|
||||
|
||||
@property
|
||||
def num_records(self) -> int:
|
||||
return self._num_records
|
||||
|
||||
def fetch_record(
|
||||
self,
|
||||
index: int,
|
||||
keys: Union[str, List[str]],
|
||||
):
|
||||
if not self._data and self._num_records == 0:
|
||||
raise RuntimeError("Store not loaded")
|
||||
if not 0 <= index < self._num_records:
|
||||
raise ValueError(
|
||||
f"Record index out of bounds: {index}, num_records={self._num_records}"
|
||||
)
|
||||
if isinstance(keys, str):
|
||||
return self._fetch_record_key(keys, index)
|
||||
return {k: self._fetch_record_key(k, index) for k in keys}
|
||||
return _record_fetch(self, index, keys)
|
||||
|
||||
def _fetch_record_key(self, key: str, index: int):
|
||||
offsets = self._offsets.get(key)
|
||||
if offsets:
|
||||
start = offsets[index]
|
||||
end = (
|
||||
offsets[index + 1]
|
||||
if index + 1 < len(offsets)
|
||||
else self._data[key][0].shape[0]
|
||||
)
|
||||
return self._data[key][0][start:end]
|
||||
return self._data[key][index]
|
||||
|
||||
def _record_fetch(self, index: int, keys: Union[str, List[str]]):
|
||||
if not getattr(self, "_data", None) and self._num_records == 0:
|
||||
raise RuntimeError("Store not loaded")
|
||||
if not 0 <= index < self._num_records:
|
||||
raise ValueError(
|
||||
f"Record index out of bounds: {index}, num_records={self._num_records}"
|
||||
)
|
||||
if isinstance(keys, str):
|
||||
return _fetch_record_key(self, keys, index)
|
||||
return {k: _fetch_record_key(self, k, index) for k in keys}
|
||||
|
||||
|
||||
def _fetch_record_key(self, key: str, index: int):
|
||||
offsets = self._offsets.get(key)
|
||||
if offsets:
|
||||
start = offsets[index]
|
||||
end = (
|
||||
offsets[index + 1]
|
||||
if index + 1 < len(offsets)
|
||||
else self._data[key][0].shape[0]
|
||||
)
|
||||
return self._data[key][0][start:end]
|
||||
return self._data[key][index]
|
||||
|
||||
|
||||
class StoreFactory(BaseFactory["Store"]):
|
||||
@@ -295,18 +421,22 @@ class H5Store(Store, Streamable, Recordable):
|
||||
Each key is stored as a group of per-record datasets (``data_0``,
|
||||
``data_1``, …). Supports both access modes:
|
||||
|
||||
- **Stream** (``fetch(begin, end, key)``): concatenates across
|
||||
records via ``_cum`` — used by SEQ/SFT where data is a token stream.
|
||||
- **Record** (``fetch_record(i, key)``): indexes ``_data[key]``
|
||||
directly — used by DPO/GRPO where each record is independent.
|
||||
|
||||
``len(store)`` returns the **token count** (stream semantics) so
|
||||
SEQ/SFT windowing works. Record-only code uses
|
||||
``store.num_records`` instead.
|
||||
- **Stream**: ``fetch(begin, end, key)`` and ``store[i]`` slice
|
||||
across concatenated records via ``_cum`` — used by SEQ/SFT.
|
||||
- **Record**: ``fetch_record(i, key)`` and ``store[i]`` (when
|
||||
``window_size == 0``) index ``_data[key]`` directly — used by
|
||||
DPO/GRPO.
|
||||
"""
|
||||
|
||||
segments_are_records = True
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
window_size: int = 0,
|
||||
stride: Optional[int] = None,
|
||||
):
|
||||
super().__init__(window_size=window_size, stride=stride)
|
||||
|
||||
def load(self, path: str, **kwargs):
|
||||
self._normalize(load_h5(path))
|
||||
|
||||
@@ -321,11 +451,12 @@ class MmapStore(Store, Streamable, Recordable):
|
||||
|
||||
Supports both access modes:
|
||||
|
||||
- **Stream** (``fetch(begin, end, key)``): always available.
|
||||
- **Stream**: always available via :meth:`fetch`.
|
||||
- **Record** (``fetch_record(i, key)``): only when ``meta.json``
|
||||
contains per-record ``offsets`` (written via
|
||||
``save_bin(..., record_keys=...)``). Legacy bin files without
|
||||
offsets have ``num_records == 0``.
|
||||
offsets have ``num_records == 0`` and ``len(store)`` reflects the
|
||||
windowed sample count when ``window_size > 0``.
|
||||
|
||||
``segments_are_records`` is ``False`` here (bin segments are
|
||||
contiguous streams, not per-record) — record access is driven
|
||||
@@ -334,6 +465,14 @@ class MmapStore(Store, Streamable, Recordable):
|
||||
|
||||
segments_are_records = False
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
window_size: int = 0,
|
||||
stride: Optional[int] = None,
|
||||
):
|
||||
super().__init__(window_size=window_size, stride=stride)
|
||||
self._mmap_refs: List[Tensor] = []
|
||||
|
||||
def load(self, path: str, **kwargs):
|
||||
self._mmap_refs = []
|
||||
root = Path(path)
|
||||
@@ -416,32 +555,26 @@ class JsonlStore(Store, Streamable, Recordable):
|
||||
|
||||
- **Eager** (default): applies a :class:`TokenizeTransform` to every
|
||||
record at load time and registers per-key tensors via
|
||||
``_normalize``. Both stream (``fetch``) and record
|
||||
(``fetch_record``) access work — stream concatenates across
|
||||
records via ``_cum``, record indexes directly.
|
||||
- **Lazy** (``processor=fn`` given): keeps raw records and defers
|
||||
tokenisation to ``fetch_record``. Only record access works.
|
||||
Used by DPO/GRPO where each record is independent.
|
||||
|
||||
``len(store)`` returns the **token count** (stream semantics) in
|
||||
eager mode so SEQ/SFT windowing works; returns the **record count**
|
||||
in lazy mode where stream access is unavailable.
|
||||
``_normalize``. Both ``fetch`` (stream) and ``fetch_record``
|
||||
(record) work.
|
||||
- **Lazy** (``processor=fn`` passed): keeps raw records and defers
|
||||
tokenisation to ``fetch_record``. Only record access works —
|
||||
``len(store)`` returns ``num_records``; stream primitives raise.
|
||||
"""
|
||||
|
||||
CONFIG_NAME = "dataset_config.json"
|
||||
segments_are_records = True
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
def __init__(
|
||||
self,
|
||||
window_size: int = 0,
|
||||
stride: Optional[int] = None,
|
||||
):
|
||||
super().__init__(window_size=window_size, stride=stride)
|
||||
self._source: Optional[JsonlSource] = None
|
||||
self._processor: Optional[Callable[[dict], Dict[str, Tensor]]] = None
|
||||
self._keys_cache: Optional[List[str]] = None
|
||||
|
||||
def __len__(self) -> int:
|
||||
if self._processor is not None:
|
||||
return self._num_records
|
||||
return self._length
|
||||
|
||||
def load(self, path: str, transform=None, processor=None, **kwargs):
|
||||
self._source = JsonlSource(path)
|
||||
records = self._source.load()
|
||||
@@ -487,4 +620,17 @@ class JsonlStore(Store, Streamable, Recordable):
|
||||
if isinstance(keys, str):
|
||||
return data[keys]
|
||||
return {k: data[k] for k in keys}
|
||||
return super().fetch_record(index, keys)
|
||||
return _record_fetch(self, index, keys)
|
||||
|
||||
def fetch(self, begin: int, end: int, keys: Union[str, List[str]]):
|
||||
if self._processor is not None:
|
||||
raise RuntimeError(
|
||||
"JsonlStore in lazy (processor) mode does not support "
|
||||
"stream fetch(); use fetch_record() instead."
|
||||
)
|
||||
return _stream_fetch(self, begin, end, keys)
|
||||
|
||||
def __getitem__(self, index: int) -> Dict[str, Tensor]:
|
||||
if self._processor is not None:
|
||||
return self.fetch_record(index, self._record_keys())
|
||||
return super().__getitem__(index)
|
||||
|
||||
Reference in New Issue
Block a user