ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor: remove H5 storage backend in favor of mmap bin

Browse Source 
- Remove H5Store, H5Writer, save_h5/load_h5 and h5py dependency
- MmapStore (bin) is the sole pre-tokenized storage backend
- Move setup_logging after imports to fix E402 in __init__.py
- Clean up unused imports across test files
- Move inline test imports to file top
This commit is contained in:
ViperEkura
2026-07-29 12:50:27 +08:00
parent c2b04d8458
commit 115192c67c
17 changed files with 91 additions and 209 deletions
Download Patch File Download Diff File Expand all files Collapse all files
astrai/__init__.py
+24 -24
View File
@@ -4,30 +4,6 @@ __author__ = "ViperEkura"
import logging
import os
def setup_logging(level: str = "INFO"):
"""Attach a handler to the ``astrai`` logger (only, not root).
Call once per process, e.g. at the top of CLI scripts.
Set ``ASTR_LOG_LEVEL`` to override the default ``INFO``.
"""
_logger = logging.getLogger("astrai")
if _logger.handlers:
return
_level = getattr(
logging, os.environ.get("ASTR_LOG_LEVEL", level).upper(), logging.INFO
)
_logger.setLevel(_level)
_handler = logging.StreamHandler()
_handler.setFormatter(
logging.Formatter(
"%(asctime)s | %(levelname)-7s | %(name)s | %(message)s",
datefmt="%Y-%m-%d %H:%M:%S",
)
)
_logger.addHandler(_handler)
from astrai.config import (
AutoRegressiveLMConfig,
BaseModelConfig,
@@ -80,6 +56,30 @@ from astrai.trainer import (
Trainer,
)
def setup_logging(level: str = "INFO"):
"""Attach a handler to the ``astrai`` logger (only, not root).
Call once per process, e.g. at the top of CLI scripts.
Set ``ASTR_LOG_LEVEL`` to override the default ``INFO``.
"""
_logger = logging.getLogger("astrai")
if _logger.handlers:
return
_level = getattr(
logging, os.environ.get("ASTR_LOG_LEVEL", level).upper(), logging.INFO
)
_logger.setLevel(_level)
_handler = logging.StreamHandler()
_handler.setFormatter(
logging.Formatter(
"%(asctime)s | %(levelname)-7s | %(name)s | %(message)s",
datefmt="%Y-%m-%d %H:%M:%S",
)
)
_logger.addHandler(_handler)
__all__ = [
"AutoRegressiveLM",
"AutoRegressiveLMConfig",
astrai/dataset/__init__.py
-6
View File
@@ -6,7 +6,6 @@ from astrai.dataset.dataset import (
)
from astrai.dataset.sampler import RDSampler
from astrai.dataset.storage import (
H5Store,
JsonlStore,
MmapStore,
Recordable,
@@ -17,9 +16,7 @@ from astrai.dataset.storage import (
)
from astrai.serialization import (
load_bin,
load_h5,
save_bin,
save_h5,
)
__all__ = [
@@ -31,12 +28,9 @@ __all__ = [
"Streamable",
"Recordable",
"StoreFactory",
"H5Store",
"MmapStore",
"JsonlStore",
"detect_format",
"save_h5",
"load_h5",
"save_bin",
"load_bin",
"RDSampler",
astrai/dataset/dataset.py
+3 -3
View File
@@ -314,7 +314,7 @@ class DatasetFactory(BaseFactory["BaseDataset"]):
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
storage_type: Storage backend ("bin", "jsonl") or
None for auto-detection.
tokenizer_path: Path to tokenizer for lazy JSONL
tokenisation (record datasets only).
@@ -384,7 +384,7 @@ class DatasetFactory(BaseFactory["BaseDataset"]):
"""Build an on-the-fly tokenisation processor if applicable.
Only raw JSONL + record datasets (DPO/GRPO) need a processor;
pre-tokenised backends (H5/bin) and stream datasets (SEQ/SFT)
pre-tokenised backends (bin) and stream datasets (SEQ/SFT)
return ``None`` so no tokenizer is loaded.
"""
if tokenizer_path is None or storage_type != "jsonl":
@@ -451,7 +451,7 @@ class DPODataset(BaseDataset):
Two loading paths (handled by :class:`DatasetFactory`):
- **Pre-tokenized** (H5/bin): ``store.load(path)`` reads per-record
- **Pre-tokenized** (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,
astrai/dataset/storage.py
+7 -45
View File
@@ -10,7 +10,6 @@ Architecture (composition over inheritance):
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)
JsonlStore(Store, Streamable, Recordable)
@@ -36,9 +35,9 @@ 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/
JSONL) or opaque shards (bin). Record access for bin relies on
``_offsets`` instead.
tells ``_normalize`` whether segments are inherently per-record (JSONL)
or opaque shards (bin). Record access for bin relies on ``_offsets``
instead.
:class:`JsonlStore` supports a lazy mode (``processor=fn``) that keeps
raw records and defers tokenisation to ``fetch_record`` — used by DPO
@@ -62,7 +61,6 @@ from astrai.preprocessing.transform import TokenizeTransform
from astrai.serialization import (
load_bin,
load_bin_offsets,
load_h5,
)
logger = logging.getLogger(__name__)
@@ -83,19 +81,10 @@ def detect_format(load_path: str) -> str:
root = Path(load_path)
if root.is_file():
suffix = root.suffix.lower()
if suffix in (".h5", ".hdf5"):
return "h5"
if suffix == ".jsonl":
return "jsonl"
raise ValueError(f"Unsupported file format: {suffix}")
h5_files = [
Path(p)
for pattern in ("*.h5", "*.hdf5")
for p in glob.glob(str(root / "**" / pattern), recursive=True)
]
if h5_files:
return "h5"
bin_files = [Path(p) for p in glob.glob(str(root / "**" / "*.bin"), recursive=True)]
if bin_files:
has_meta = (root / "meta.json").exists() or len(
@@ -185,7 +174,7 @@ class Store(ABC):
"""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``).
indexing (JSONL segments or bin ``_offsets``).
"""
return self._num_records
@@ -269,7 +258,7 @@ class Store(ABC):
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
``segments_are_records`` is True (JSONL), ``_data[key]`` is
a per-record list and ``fetch_record`` indexes it directly.
Nested keys (GRPO ``responses``/``masks`` as
@@ -305,7 +294,7 @@ class Store(ABC):
logger.warning(
"Key '%s' has %d segments with offsets — record mode "
"disabled for this key (multi-shard bin+offsets not "
"supported). Merge shards or use H5/JSONL.",
"supported). Merge shards or use JSONL.",
key,
len(segs),
)
@@ -330,7 +319,7 @@ class Streamable:
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
that can also serve record access (JSONL/bin+offsets), the
``fetch_record`` API from :class:`Recordable` is used instead.
"""
@@ -415,33 +404,6 @@ class StoreFactory(BaseFactory["Store"]):
"""Factory for creating Store instances by type name."""
@StoreFactory.register("h5")
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``,
``data_1``, …). Supports both access modes:
- **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))
@StoreFactory.register("bin")
class MmapStore(Store, Streamable, Recordable):
"""Memory-mapped binary storage backend.
astrai/preprocessing/pipeline.py
+1 -1
View File
@@ -1,7 +1,7 @@
"""Config-driven JSONL preprocessing pipeline.
Composes a :class:`BaseMaskBuilder` (selected by ``input.type``) with
sharding and flush to ``.h5`` / ``.bin`` storage. Packing, position-id
sharding and flush to ``.bin`` storage. Packing, position-id
generation and storage writing are each delegated to pluggable strategies,
dispatched by configuration keys.
astrai/preprocessing/writer.py
+2 -21
View File
@@ -1,7 +1,7 @@
"""Storage writer strategies for pipeline output.
The :class:`StoreWriter` abstraction decouples the pipeline from the
concrete storage format (bin / h5). The pipeline builds a ``{key:
concrete storage format (bin). The pipeline builds a ``{key:
List[Tensor]}`` dict and delegates the write to the writer selected
by ``output.storage_format``.
"""
@@ -15,7 +15,7 @@ from typing import Dict, List
import torch
from astrai.factory import BaseFactory
from astrai.serialization import save_bin, save_h5
from astrai.serialization import save_bin
logger = logging.getLogger(__name__)
@@ -54,22 +54,3 @@ class BinWriter(StoreWriter):
exc_info=True,
)
raise
@StoreWriterFactory.register("h5")
class H5Writer(StoreWriter):
def save(self, output_dir, domain, shard_idx, tensors):
chunk_dir = os.path.join(output_dir, domain)
file_path = os.path.join(chunk_dir, f"data_{shard_idx:04d}.h5")
try:
save_h5(chunk_dir, f"data_{shard_idx:04d}", tensors)
except Exception:
if os.path.exists(file_path):
os.remove(file_path)
logger.error(
"Failed to write shard %s/data_%04d.h5, cleaned up partial output",
domain,
shard_idx,
exc_info=True,
)
raise
astrai/serialization/__init__.py
-4
View File
@@ -20,9 +20,7 @@ from astrai.serialization.checkpoint import (
from astrai.serialization.dataset import (
load_bin,
load_bin_offsets,
load_h5,
save_bin,
save_h5,
)
__all__ = [
@@ -39,7 +37,5 @@ __all__ = [
"save_torch",
"load_bin",
"load_bin_offsets",
"load_h5",
"save_bin",
"save_h5",
]
astrai/serialization/dataset.py
+4 -45
View File
@@ -1,55 +1,14 @@
"""Dataset storage serialization helpers (HDF5 / memory-mapped binary)."""
"""Dataset storage serialization helpers (memory-mapped binary)."""
import json
import os
from pathlib import Path
from typing import Any, Dict, List, Optional
import h5py
import numpy as np
import torch
from torch import Tensor
def save_h5(file_path: str, file_name: str, tensor_group: Dict[str, List[Tensor]]):
os.makedirs(file_path, exist_ok=True)
full_file_path = os.path.join(file_path, f"{file_name}.h5")
with h5py.File(full_file_path, "w") as f:
for key, tensors in tensor_group.items():
grp = f.create_group(key)
for idx, tensor in enumerate(tensors):
arr = tensor.cpu().numpy()
grp.create_dataset(f"data_{idx}", data=arr)
def load_h5(file_path: str, share_memory=True) -> Dict[str, List[Tensor]]:
tensor_group: Dict[str, List[Tensor]] = {}
root_path = Path(file_path)
if root_path.is_file() and root_path.suffix in (".h5", ".hdf5"):
h5_files = [root_path]
else:
h5_files = list(root_path.rglob("*.h5")) + list(root_path.rglob("*.hdf5"))
for h5_file in h5_files:
with h5py.File(h5_file, "r") as f:
for key in f.keys():
grp = f[key]
dsets = []
for dset_name in grp.keys():
dset = grp[dset_name]
tensor = torch.from_numpy(dset[:])
if share_memory:
tensor = tensor.share_memory_()
dsets.append(tensor)
if tensor_group.get(key) is None:
tensor_group[key] = []
tensor_group[key].extend(dsets)
return tensor_group
def save_bin(
file_path: str,
tensor_group: Dict[str, List[Tensor]],
@@ -65,7 +24,7 @@ def save_bin(
offsets, preserving backward compatibility.
Nested keys (``List[List[Tensor]]`` such as GRPO ``responses``) are
not supported in bin format — use H5 for those.
not supported in bin format — use JSONL for those.
"""
os.makedirs(file_path, exist_ok=True)
record_keys = set(record_keys or [])
@@ -74,7 +33,7 @@ def save_bin(
if tensors and isinstance(tensors[0], list):
raise ValueError(
f"Nested key '{key}' (List[List[Tensor]]) is not supported "
f"in bin format. Use H5 or JSONL storage instead."
f"in bin format. Use JSONL storage instead."
)
cat = torch.cat(tensors, dim=0)
entry: Dict[str, Any] = {
@@ -112,7 +71,7 @@ def load_bin_offsets(file_path: str) -> Dict[str, List[int]]:
Returns an empty dict when no key has offsets (legacy bin files),
in which case record-mode access falls back to per-record segment
indexing (H5/JSONL layout).
indexing (JSONL layout).
"""
with open(os.path.join(file_path, "meta.json"), "r") as f:
meta = json.load(f)
pyproject.toml
-1
View File
@@ -8,7 +8,6 @@ name = "astrai"
readme = "README.md"
requires-python = ">=3.12"
dependencies = [
"h5py==3.15.1",
"numpy==2.4.4",
"torch==2.11.0",
"tokenizers==0.21.4",
scripts/tools/generate.py
+1 -3
View File
@@ -145,9 +145,7 @@ def processor(
@click.option("--top_p", type=float, default=0.95, help="Top-p filtering.")
@click.option("--batch_size", type=int, default=1, help="Batch size.")
@click.option("--num_samples", type=int, default=1, help="Responses per prompt.")
@click.option(
"--max_seq_len", type=int, default=2048, help="KV cache length."
)
@click.option("--max_seq_len", type=int, default=2048, help="KV cache length.")
@click.option("--frequency_penalty", type=float, default=0.0, help="Frequency penalty.")
@click.option(
"--rep_window", type=int, default=64, help="Window size for frequency penalty."
scripts/tools/preprocess.py
+3 -3
View File
@@ -1,4 +1,4 @@
"""CLI: JSONL → tokenized .h5/.bin via config-driven Pipeline."""
"""CLI: JSONL → tokenized .bin via config-driven Pipeline."""
import click
@@ -8,7 +8,7 @@ from astrai.preprocessing.pipeline import Pipeline
@click.command(
name="preprocess", help="Tokenize and pack raw JSONL data into .bin/.h5 format."
name="preprocess", help="Tokenize and pack raw JSONL data into .bin format."
)
@click.argument("inputs", nargs=-1, type=click.Path(exists=True), required=True)
@click.option(
@@ -30,7 +30,7 @@ from astrai.preprocessing.pipeline import Pipeline
)
@click.option("--batch_size", type=int, default=None, help="Records per batch.")
def preprocess_command(inputs, output_dir, pipeline_config, tokenizer_path, batch_size):
"""Tokenize and pack raw JSONL data into .bin/.h5 format."""
"""Tokenize and pack raw JSONL data into .bin format."""
config = PipelineConfig.from_file(pipeline_config)
if batch_size is not None:
if batch_size < 1:
tests/data/test_dataset.py
+42 -46
View File
@@ -7,18 +7,24 @@ import pytest
import torch
from astrai.config.preprocess_config import PipelineConfig
from astrai.dataset.dataset import DatasetFactory, dpo_tokenize
from astrai.dataset.dataset import (
DatasetFactory,
GRPODataset,
dpo_tokenize,
grpo_collate_fn,
)
from astrai.dataset.storage import (
H5Store,
JsonlStore,
MmapStore,
StoreFactory,
detect_format,
)
from astrai.preprocessing.builder import SectionedMaskBuilder
from astrai.serialization import (
load_bin,
save_bin,
save_h5,
)
from tests.data.conftest import make_grpo_no_template_config
def _rand_seq(length, vocab=1000):
@@ -70,7 +76,7 @@ def _make_seq_dataset(
):
if data is None:
data = {"sequence": [_rand_seq(seq_length)]}
save_h5(test_dir, name, data)
save_bin(test_dir, data)
return DatasetFactory.load(
train_type,
test_dir,
@@ -83,12 +89,15 @@ def test_dataset_loader_random_paths(base_test_env):
"""Test dataset loader with multiple random paths"""
test_dir = base_test_env["test_dir"]
loaded_dataset = None
num_files = np.random.randint(2, 5)
for i in range(num_files):
seq_length = np.random.randint(200, 400)
dummy_data = {"sequence": [_rand_seq(seq_length) for _ in range(10)]}
sub_dir = os.path.join(test_dir, f"sub_{i}")
os.makedirs(sub_dir, exist_ok=True)
loaded_dataset = _make_seq_dataset(
test_dir, f"data_{i}", seq_length, data=dummy_data
sub_dir, f"data_{i}", seq_length, data=dummy_data
)
assert loaded_dataset is not None
assert len(loaded_dataset) > 0
@@ -113,8 +122,15 @@ def test_dpo_strategy_with_random_data(base_test_env):
"chosen_mask": [torch.ones(seq_length, dtype=torch.bool)],
"rejected_mask": [torch.ones(seq_length, dtype=torch.bool)],
}
dpo_dataset = _make_seq_dataset(
test_dir, "dpo_data", seq_length, train_type="dpo", data=dummy_data
save_bin(
test_dir,
dummy_data,
record_keys=["chosen", "rejected", "chosen_mask", "rejected_mask"],
)
dpo_dataset = DatasetFactory.load(
train_type="dpo",
load_path=test_dir,
window_size=0,
)
assert dpo_dataset is not None
@@ -196,24 +212,20 @@ def test_dataset_too_short_for_window(base_test_env):
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)
store = MmapStore(window_size=64, stride=64)
with pytest.raises(IndexError, match="Data too short"):
store.sample_window(0)
def test_unloaded_dataset_len():
"""__len__ on a store with no data returns 0."""
from astrai.dataset.storage import H5Store
store = H5Store(window_size=64, stride=64)
store = MmapStore(window_size=64, stride=64)
assert len(store) == 0
def test_store_unloaded_len():
"""Unloaded Store has __len__ == 0"""
store = H5Store()
store = MmapStore()
assert len(store) == 0
assert store.keys == []
@@ -227,7 +239,7 @@ def test_store_fetch_begin_equals_end(base_test_env):
def test_store_fetch_before_load():
"""Store.fetch before load raises RuntimeError"""
store = H5Store()
store = MmapStore()
with pytest.raises(RuntimeError, match="not loaded"):
store.fetch(0, 10, "sequence")
@@ -255,9 +267,7 @@ def test_create_store_invalid_type():
def test_store_multi_segment_concat(base_test_env):
"""Multi-segment H5 data is concatenated into single tensor at load time"""
import os
"""Multi-segment data is concatenated into single tensor at load time"""
test_dir = base_test_env["test_dir"]
data_dir = os.path.join(test_dir, "multi_seg")
os.makedirs(data_dir, exist_ok=True)
@@ -267,9 +277,9 @@ def test_store_multi_segment_concat(base_test_env):
torch.tensor([4, 5, 6, 7]),
torch.tensor([8, 9]),
]
save_h5(data_dir, "data", {"sequence": segs})
save_bin(data_dir, {"sequence": segs})
store = StoreFactory.create("h5")
store = StoreFactory.create("bin")
store.load(data_dir)
assert store.token_count == 9
result = store.fetch(2, 7, "sequence")
@@ -321,7 +331,7 @@ def test_mmap_dataset_load(base_test_env):
def test_normalize_empty_key():
"""_normalize with empty tensor list does not crash."""
store = H5Store()
store = MmapStore()
store._normalize({"sequence": []})
assert len(store) == 0
assert store.num_records == 0 # empty key forces num_records=0
@@ -330,7 +340,7 @@ def test_normalize_empty_key():
def test_normalize_mixed_empty_key():
"""_normalize with empty + non-empty keys returns min=0 records."""
store = H5Store()
store = MmapStore()
store._normalize({"sequence": [torch.tensor([1, 2, 3])], "loss_mask": []})
assert len(store) == 0
assert store.num_records == 0
@@ -340,8 +350,6 @@ def test_normalize_mixed_empty_key():
def test_grpo_dataset_dtype(base_test_env):
"""GRPO dataset returns correct dtypes for per-record structured data."""
from astrai.dataset.dataset import GRPODataset
G = 4
store = type(
"FakeStore",
@@ -373,8 +381,6 @@ def test_grpo_dataset_dtype(base_test_env):
def test_grpo_dataset_load(base_test_env):
"""GRPO dataset loads record-structured data with per-response boundaries."""
from astrai.dataset.dataset import GRPODataset
G = 3
prompt_len = 8
resp_lens = [5, 7, 4]
@@ -430,15 +436,14 @@ def test_detect_format_bin_dir(base_test_env):
def test_store_fetch_multi_key(base_test_env):
test_dir = base_test_env["test_dir"]
save_h5(
save_bin(
test_dir,
"multi_key",
{
"sequence": [torch.randint(0, 100, (100,), dtype=torch.int64)],
"loss_mask": [torch.ones(100, dtype=torch.int64)],
},
)
store = StoreFactory.create("h5")
store = StoreFactory.create("bin")
store.load(test_dir)
result = store.fetch(10, 20, ["sequence", "loss_mask"])
assert isinstance(result, dict)
@@ -448,8 +453,8 @@ def test_store_fetch_multi_key(base_test_env):
def test_store_fetch_out_of_bounds(base_test_env):
test_dir = base_test_env["test_dir"]
save_h5(test_dir, "bounds", {"sequence": [torch.randint(0, 100, (50,))]})
store = StoreFactory.create("h5")
save_bin(test_dir, {"sequence": [torch.randint(0, 100, (50,))]})
store = StoreFactory.create("bin")
store.load(test_dir)
with pytest.raises(ValueError, match="out of bounds"):
store.fetch(-1, 10, "sequence")
@@ -461,7 +466,7 @@ def test_store_fetch_out_of_bounds(base_test_env):
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")
dataset = _make_seq_dataset(test_dir, "explicit", storage_type="bin")
assert len(dataset) > 0
assert dataset.token_count == 200
@@ -820,9 +825,6 @@ def _write_grpo_jsonl(test_dir, tokenizer_path, records):
def test_grpo_builder_preserves_response_boundaries(base_test_env):
"""MultiOutputMaskBuilder with list_field returns List[List[int]] for responses."""
from astrai.preprocessing.builder import SectionedMaskBuilder
from tests.data.conftest import make_grpo_no_template_config
tokenizer = base_test_env["tokenizer"]
_save_test_tokenizer(base_test_env["test_dir"], tokenizer)
@@ -863,8 +865,6 @@ def test_grpo_builder_preserves_response_boundaries(base_test_env):
def test_grpo_end_to_end_jsonl(base_test_env):
"""Full GRPO pipeline: JSONL → JsonlStore → GRPODataset → collate_fn."""
from astrai.dataset.dataset import grpo_collate_fn
test_dir = base_test_env["test_dir"]
tokenizer = base_test_env["tokenizer"]
tokenizer_path = _save_test_tokenizer(test_dir, tokenizer)
@@ -913,8 +913,6 @@ def test_grpo_end_to_end_jsonl(base_test_env):
def test_grpo_collate_variable_lengths():
"""collate_fn pads variable-length responses to [B, G, R_max]."""
from astrai.dataset.dataset import grpo_collate_fn
batch = [
{
"prompts": torch.tensor([1, 2, 3]),
@@ -954,8 +952,6 @@ def test_grpo_collate_variable_lengths():
def test_grpo_multiple_records(base_test_env):
"""GRPODataset loads multiple records with correct structure."""
from astrai.dataset.dataset import GRPODataset
G = 4
n_records = 5
@@ -1128,8 +1124,8 @@ def test_jsonl_store_eager_len_returns_token_count(base_test_env):
assert len(store.keys) > 0
def test_h5_store_dual_mode(base_test_env):
"""H5Store supports both fetch (stream) and fetch_record (record).
def test_mmap_store_dual_mode(base_test_env):
"""MmapStore 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
@@ -1143,9 +1139,9 @@ def test_h5_store_dual_mode(base_test_env):
"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)
save_bin(test_dir, dummy_data, record_keys=["chosen", "rejected"])
store = H5Store()
store = MmapStore()
store.load(test_dir)
assert store.token_count == seq_length * 2
@@ -1160,7 +1156,7 @@ def test_h5_store_dual_mode(base_test_env):
# 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 = MmapStore(window_size=seq_length, stride=seq_length)
stream_view.load(test_dir)
assert len(stream_view) == 1
tests/data/test_preprocess_config.py
+2 -2
View File
@@ -30,7 +30,7 @@ def test_from_dict_flat():
"mask": {"system": "mask", "assistant": "train"},
"mask_default": "mask",
"preprocessing": {"max_seq_len": 1024},
"output": {"storage_format": "h5"},
"output": {"storage_format": "bin"},
}
config = PipelineConfig.from_dict(data)
assert config.input.sections == [
@@ -38,7 +38,7 @@ def test_from_dict_flat():
]
assert config.mask == {"system": "mask", "assistant": "train"}
assert config.preprocessing.max_seq_len == 1024
assert config.output.storage_format == "h5"
assert config.output.storage_format == "bin"
def test_to_dict_roundtrip():
tests/inference/test_task.py
+1 -1
View File
@@ -2,7 +2,7 @@
from unittest.mock import MagicMock
from astrai.inference import STOP, Task, TaskManager, TaskStatus
from astrai.inference import Task, TaskManager, TaskStatus
def _make_mock_tokenizer():
tests/module/test_tie_weight.py
-1
View File
@@ -1,7 +1,6 @@
import json
import os
import pytest
import safetensors.torch as st
import torch
tests/trainer/conftest.py
-2
View File
@@ -1,5 +1,3 @@
import os
import pytest
import torch
tests/trainer/test_grpo_strategy.py
+1 -1
View File
@@ -3,7 +3,7 @@ import torch
from astrai.model.transformer import AutoRegressiveLM
from astrai.trainer.strategy import GRPOStrategy
from tests.helpers import FakeExecutor, make_frozen, make_model, make_rollout_config
from tests.helpers import FakeExecutor, make_frozen, make_model
def _make_batch(