From 553a42702d927b09710c4bb62f5f9455fd24461b Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Sat, 18 Jul 2026 23:20:41 +0800 Subject: [PATCH] refactor: replace diamond inheritance with mixin composition MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - StreamStore/RecordStore → Streamable/Recordable (stateless mixins) - Store is sole base class, no MRO ambiguity - H5Store/MmapStore/JsonlStore mix in both traits explicitly - segments_are_records declared per-subclass (H5/Jsonl=True, bin=False) - Add tests for dpo_tokenize, lazy jsonl, dual-mode H5, stream-only bin - Remove unused _to_tensor helper --- astrai/dataset/__init__.py | 8 +- astrai/dataset/dataset.py | 30 ++++--- astrai/dataset/storage.py | 93 ++++++++++++++------- tests/data/test_dataset.py | 167 ++++++++++++++++++++++++++++++++++++- 4 files changed, 251 insertions(+), 47 deletions(-) diff --git a/astrai/dataset/__init__.py b/astrai/dataset/__init__.py index c5c45c9..6533708 100644 --- a/astrai/dataset/__init__.py +++ b/astrai/dataset/__init__.py @@ -9,10 +9,10 @@ from astrai.dataset.storage import ( H5Store, JsonlStore, MmapStore, - RecordStore, + Recordable, Store, StoreFactory, - StreamStore, + Streamable, detect_format, ) from astrai.serialization import ( @@ -28,8 +28,8 @@ __all__ = [ "dpo_collate_fn", "grpo_collate_fn", "Store", - "StreamStore", - "RecordStore", + "Streamable", + "Recordable", "StoreFactory", "H5Store", "MmapStore", diff --git a/astrai/dataset/dataset.py b/astrai/dataset/dataset.py index aac2ed4..d32c866 100644 --- a/astrai/dataset/dataset.py +++ b/astrai/dataset/dataset.py @@ -1,4 +1,22 @@ -"""Dataset implementations with factory pattern for training.""" +"""Dataset implementations with factory pattern for training. + +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 + +``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. + +``__len__`` returns the sample count (stream: windows, record: +records) so DataLoader and progress bars work uniformly. +""" from abc import ABC, abstractmethod from functools import partial @@ -9,24 +27,14 @@ from torch import Tensor from torch.utils.data import Dataset from astrai.dataset.storage import ( - RecordStore, Store, StoreFactory, - StreamStore, detect_format, ) from astrai.factory import BaseFactory from astrai.tokenize import AutoTokenizer -def _to_tensor(value: list, dtype: Optional[torch.dtype] = None) -> Tensor: - if dtype is not None: - return torch.tensor(value, dtype=dtype) - if value and isinstance(value[0], bool): - return torch.tensor(value, dtype=torch.bool) - return torch.tensor(value, dtype=torch.int32) - - def dpo_tokenize( record: dict, tokenizer, diff --git a/astrai/dataset/storage.py b/astrai/dataset/storage.py index 2047a00..c0d09a3 100644 --- a/astrai/dataset/storage.py +++ b/astrai/dataset/storage.py @@ -1,24 +1,37 @@ """Storage backends for different data formats. -Two access modes are reflected in the class hierarchy: +Architecture (mixin composition, no diamond inheritance): -- :class:`StreamStore` — ``fetch(begin, end, key)`` slices across - concatenated segments. Used by PT/SFT where data is a long token - stream. ``len(store)`` returns the total token count. -- :class:`RecordStore` — ``fetch_record(i, key)`` returns the *i*-th - record without cross-record concatenation. Used by DPO/GRPO where - each record is an independent training unit. ``num_records`` returns + 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 + + H5Store(Store, Streamable, Recordable) + MmapStore(Store, Streamable, Recordable) + JsonlStore(Store, Streamable, Recordable) + +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. + +Access-mode semantics: + +- **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. -Both share ``_data`` / ``_cum`` / ``_offsets`` bookkeeping via the -common :class:`Store` base, which also owns ``_normalize`` for -registering segments. Subclasses pick the access mode by inheriting -from the appropriate base. +``segments_are_records`` (class attribute on each Store subclass) +tells ``_normalize`` whether segments are inherently per-record (H5/ +JSONL) or opaque shards (bin). Record access for bin relies on +``_offsets`` instead. -:class:`ProcessedStore` composes a :class:`JsonlSource` (raw record -reader) with a pure ``record -> dict_of_tensors`` processor so that -DPO/GRPO can tokenise raw JSONL on the fly without a pre-tokenised -H5/bin file. This keeps the tokenizer out of the Store base. +:class:`JsonlStore` supports a lazy mode (``processor=fn``) that keeps +raw records and defers tokenisation to ``fetch_record`` — used by DPO +to train directly from a ``.jsonl`` file without a pre-tokenised copy. """ import bisect @@ -112,6 +125,14 @@ class Store(ABC): def keys(self) -> List[str]: return list(self._data.keys()) + def __len__(self) -> int: + """Default: token count (stream semantics). + + Subclasses that are record-only (e.g. lazy JsonlStore) override + to return ``self._num_records``. + """ + return self._length + def _normalize( self, raw: Dict[str, list], @@ -181,11 +202,13 @@ class Store(ABC): self._num_records = 0 -class StreamStore(Store): - """Store exposing stream access: ``fetch(begin, end, key)``.""" +class Streamable: + """Mixin: stream access ``fetch(begin, end, key)``. - def __len__(self) -> int: - return self._length + 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. + """ def fetch( self, @@ -200,10 +223,10 @@ class StreamStore(Store): f"Index out of bounds: begin={begin}, end={end}, length={self._length}" ) if isinstance(keys, str): - return self._fetch_key(keys, begin, end) - return {k: self._fetch_key(k, begin, end) for k in keys} + return self._fetch_stream_key(keys, begin, end) + return {k: self._fetch_stream_key(k, begin, end) for k in keys} - def _fetch_key(self, key: str, begin: int, end: int) -> Tensor: + 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) @@ -219,11 +242,12 @@ class StreamStore(Store): return results[0] if len(results) == 1 else torch.cat(results, dim=0) -class RecordStore(Store): - """Mixin exposing record access: ``fetch_record(i, key)``. +class Recordable: + """Mixin: record access ``fetch_record(i, key)``. - ``__len__`` is **not** overridden — subclasses decide whether - ``len()`` returns token count (stream) or record count (record-only). + 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. """ segments_are_records = True @@ -237,7 +261,7 @@ class RecordStore(Store): index: int, keys: Union[str, List[str]], ): - if not self._data: + if not self._data and self._num_records == 0: raise RuntimeError("Store not loaded") if not 0 <= index < self._num_records: raise ValueError( @@ -265,7 +289,7 @@ class StoreFactory(BaseFactory["Store"]): @StoreFactory.register("h5") -class H5Store(StreamStore, RecordStore): +class H5Store(Store, Streamable, Recordable): """HDF5-based storage backend (pre-tokenized data). Each key is stored as a group of per-record datasets (``data_0``, @@ -281,12 +305,14 @@ class H5Store(StreamStore, RecordStore): ``store.num_records`` instead. """ + segments_are_records = True + def load(self, path: str, **kwargs): self._normalize(load_h5(path)) @StoreFactory.register("bin") -class MmapStore(StreamStore, RecordStore): +class MmapStore(Store, Streamable, Recordable): """Memory-mapped binary storage backend. Each key is a single .bin file backed by ``np.memmap(mode="r")``. @@ -301,9 +327,13 @@ class MmapStore(StreamStore, RecordStore): ``save_bin(..., record_keys=...)``). Legacy bin files without offsets have ``num_records == 0``. - ``len(store)`` returns the **token count** (stream semantics). + ``segments_are_records`` is ``False`` here (bin segments are + contiguous streams, not per-record) — record access is driven + purely by ``_offsets``. """ + segments_are_records = False + def load(self, path: str, **kwargs): self._mmap_refs = [] root = Path(path) @@ -375,7 +405,7 @@ class JsonlSource: @StoreFactory.register("jsonl") -class JsonlStore(StreamStore, RecordStore): +class JsonlStore(Store, Streamable, Recordable): """JSONL reader with two tokenisation modes. A JSONL dataset is a ``.jsonl`` file or a directory of ``*.jsonl`` @@ -399,6 +429,7 @@ class JsonlStore(StreamStore, RecordStore): """ CONFIG_NAME = "dataset_config.json" + segments_are_records = True def __init__(self): super().__init__() diff --git a/tests/data/test_dataset.py b/tests/data/test_dataset.py index be85e9f..a98ebf1 100644 --- a/tests/data/test_dataset.py +++ b/tests/data/test_dataset.py @@ -1,14 +1,16 @@ import json import os +import tempfile import numpy as np import pytest import torch from astrai.config.preprocess_config import PipelineConfig -from astrai.dataset.dataset import DatasetFactory, SEQDataset +from astrai.dataset.dataset import DatasetFactory, SEQDataset, dpo_tokenize from astrai.dataset.storage import ( H5Store, + JsonlStore, StoreFactory, detect_format, ) @@ -900,3 +902,166 @@ def test_grpo_multiple_records(base_test_env): assert item["rewards"].shape == (G,) for g in range(G): assert item["responses"][g].shape == item["masks"][g].shape + + +def _write_dpo_jsonl(test_dir, records): + """Write a raw DPO JSONL file (no dataset_config.json).""" + path = os.path.join(test_dir, "dpo.jsonl") + with open(path, "w", encoding="utf-8") as f: + for rec in records: + f.write(json.dumps(rec, ensure_ascii=False) + "\n") + return path + + +def test_dpo_tokenize_pure_function(): + """dpo_tokenize returns flat lists with correct mask alignment.""" + + class FakeTokenizer: + def encode(self, text, add_special_tokens=True): + return [len(text)] if add_special_tokens else [len(text) + 1] + + record = {"input": "ab", "chosen": "xyz", "rejected": "w"} + result = dpo_tokenize(record, FakeTokenizer(), max_len=64, pad_id=0) + + assert set(result.keys()) == {"chosen", "rejected", "chosen_mask", "rejected_mask"} + assert len(result["chosen"]) == len(result["chosen_mask"]) + assert len(result["rejected"]) == len(result["rejected_mask"]) + assert len(result["chosen"]) == len(result["rejected"]) + + assert result["chosen_mask"][0] == 0 + assert any(m == 1 for m in result["chosen_mask"]) + assert result["rejected_mask"][0] == 0 + + +def test_dpo_tokenize_malformed_record(): + """dpo_tokenize returns None for missing fields.""" + + class FakeTokenizer: + def encode(self, text, add_special_tokens=True): + return [1] + + assert dpo_tokenize({}, FakeTokenizer()) is None + assert dpo_tokenize({"input": "a"}, FakeTokenizer()) is None + assert dpo_tokenize({"input": "a", "chosen": "b"}, FakeTokenizer()) is None + + +def test_dpo_jsonl_lazy_load(base_test_env): + """DPODataset loads raw JSONL with tokenizer_path → lazy processor.""" + test_dir = base_test_env["test_dir"] + tokenizer_path = _save_test_tokenizer(test_dir, base_test_env["tokenizer"]) + + records = [ + {"input": "Hello", "chosen": "world", "rejected": "earth"}, + {"input": "Foo", "chosen": "bar", "rejected": "baz"}, + ] + path = _write_dpo_jsonl(test_dir, records) + + ds = DatasetFactory.load( + train_type="dpo", + load_path=path, + window_size=0, + tokenizer_path=tokenizer_path, + ) + + assert len(ds) == 2 + assert ds.storage.num_records == 2 + assert ds.storage._processor is not None + + item = ds[0] + assert set(item.keys()) == {"chosen", "rejected", "chosen_mask", "rejected_mask"} + assert item["chosen"].dtype == torch.long + assert item["chosen_mask"].dtype == torch.bool + assert item["chosen"].shape == item["chosen_mask"].shape + assert item["chosen"].shape == item["rejected"].shape + + +def test_dpo_jsonl_lazy_no_tokenizer(): + """DPODataset on jsonl without tokenizer_path falls back to eager + (which requires dataset_config.json, so it should raise).""" + with tempfile.TemporaryDirectory() as d: + path = os.path.join(d, "dpo.jsonl") + with open(path, "w") as f: + f.write(json.dumps({"input": "a", "chosen": "b", "rejected": "c"}) + "\n") + + with pytest.raises(FileNotFoundError, match="dataset_config.json"): + DatasetFactory.load( + train_type="dpo", + load_path=path, + window_size=0, + ) + + +def test_jsonl_store_lazy_len_returns_record_count(base_test_env): + """JsonlStore in lazy mode: len() returns record count, not tokens.""" + test_dir = base_test_env["test_dir"] + records = [{"input": str(i), "chosen": "c", "rejected": "r"} for i in range(5)] + path = _write_dpo_jsonl(test_dir, records) + + store = JsonlStore() + store.load(path, processor=lambda r: {"chosen": torch.tensor([1, 2])}) + + assert len(store) == 5 + assert store.num_records == 5 + + +def test_jsonl_store_eager_len_returns_token_count(base_test_env): + """JsonlStore in eager mode: num_records reflects per-record count.""" + test_dir = base_test_env["test_dir"] + tokenizer_path = _save_test_tokenizer(test_dir, base_test_env["tokenizer"]) + data_dir = _write_jsonl_dataset( + test_dir, + tokenizer_path, + [{"text": "hello world"}, {"text": "foo bar"}], + config_overrides={ + "preprocessing": {"max_seq_len": 128, "min_chars": 0}, + "output": {"position_ids_mode": "none"}, + }, + ) + + store = JsonlStore() + store.load(data_dir) + + assert store.num_records == 2 + assert len(store.keys) > 0 + + +def test_h5_store_dual_mode(base_test_env): + """H5Store supports both fetch (stream) and fetch_record (record).""" + test_dir = base_test_env["test_dir"] + + seq_length = 64 + dummy_data = { + "chosen": [_rand_seq(seq_length), _rand_seq(seq_length)], + "rejected": [_rand_seq(seq_length), _rand_seq(seq_length)], + } + save_h5(test_dir, "dpo_data", dummy_data) + + store = H5Store() + store.load(test_dir) + + assert len(store) == seq_length * 2 + assert store.num_records == 2 + + rec0 = store.fetch_record(0, "chosen") + assert rec0.shape == (seq_length,) + + stream = store.fetch(0, 10, "chosen") + assert stream.shape == (10,) + + +def test_mmap_store_stream_only_no_offsets(base_test_env): + """MmapStore without offsets: num_records == 0, stream works.""" + test_dir = base_test_env["test_dir"] + + seq_length = 128 + dummy_data = {"sequence": [_rand_seq(seq_length)]} + save_bin(test_dir, dummy_data) + + store = StoreFactory.create("bin") + store.load(test_dir) + + assert len(store) == seq_length + assert store.num_records == 0 + + chunk = store.fetch(0, 32, "sequence") + assert chunk.shape == (32,)