refactor: remove H5 storage backend in favor of mmap bin

- 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

astrai/__init__.py

@@ -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,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

@@ -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

@@ -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/
+tells ``_normalize`` whether segments are inherently per-record (JSONL)
-JSONL) or opaque shards (bin).  Record access for bin relies on
+or opaque shards (bin).  Record access for bin relies on ``_offsets``
-``_offsets`` instead.
+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,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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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(
+@click.option("--max_seq_len", type=int, default=2048, help="KV cache length.")
-    "--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

@@ -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

@@ -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(
+    save_bin(
-        test_dir, "dpo_data", seq_length, train_type="dpo", data=dummy_data
+        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 = MmapStore(window_size=64, stride=64)
-
-    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__ on a store with no data returns 0."""
-    from astrai.dataset.storage import H5Store
+    store = MmapStore(window_size=64, stride=64)
-
-    store = H5Store(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"""
+    """Multi-segment data is concatenated into single tensor at load time"""
-    import os
-
     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,))]})
+    save_bin(test_dir, {"sequence": [torch.randint(0, 100, (50,))]})
-    store = StoreFactory.create("h5")
+    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):
+def test_mmap_store_dual_mode(base_test_env):
-    """H5Store supports both fetch (stream) and fetch_record (record).
+    """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

@@ -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

@@ -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,7 +1,6 @@
 import json
 import os
 
-import pytest
 import safetensors.torch as st
 import torch
 

tests/trainer/conftest.py

@@ -1,5 +1,3 @@
-import os
-
 import pytest
 import torch
 

tests/trainer/test_grpo_strategy.py

@@ -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(