diff --git a/astrai/__init__.py b/astrai/__init__.py index 6c65541..c0b9ebf 100644 --- a/astrai/__init__.py +++ b/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", diff --git a/astrai/dataset/__init__.py b/astrai/dataset/__init__.py index 6533708..7431137 100644 --- a/astrai/dataset/__init__.py +++ b/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", diff --git a/astrai/dataset/dataset.py b/astrai/dataset/dataset.py index 2dcb1d4..6bcff35 100644 --- a/astrai/dataset/dataset.py +++ b/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, diff --git a/astrai/dataset/storage.py b/astrai/dataset/storage.py index fb4522e..698f913 100644 --- a/astrai/dataset/storage.py +++ b/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/ -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. diff --git a/astrai/preprocessing/pipeline.py b/astrai/preprocessing/pipeline.py index 009bd41..012e12a 100644 --- a/astrai/preprocessing/pipeline.py +++ b/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. diff --git a/astrai/preprocessing/writer.py b/astrai/preprocessing/writer.py index 42488a6..de1d98a 100644 --- a/astrai/preprocessing/writer.py +++ b/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 diff --git a/astrai/serialization/__init__.py b/astrai/serialization/__init__.py index 6dd11c5..501f886 100644 --- a/astrai/serialization/__init__.py +++ b/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", ] diff --git a/astrai/serialization/dataset.py b/astrai/serialization/dataset.py index 790f5fa..4b69d18 100644 --- a/astrai/serialization/dataset.py +++ b/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) diff --git a/pyproject.toml b/pyproject.toml index 0f2ab09..2da7f86 100644 --- a/pyproject.toml +++ b/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", diff --git a/scripts/tools/generate.py b/scripts/tools/generate.py index 3e498e9..9263878 100644 --- a/scripts/tools/generate.py +++ b/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( - "--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." diff --git a/scripts/tools/preprocess.py b/scripts/tools/preprocess.py index 3e4e4b2..1b3b364 100644 --- a/scripts/tools/preprocess.py +++ b/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: diff --git a/tests/data/test_dataset.py b/tests/data/test_dataset.py index d04b904..90f5e07 100644 --- a/tests/data/test_dataset.py +++ b/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( - 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 diff --git a/tests/data/test_preprocess_config.py b/tests/data/test_preprocess_config.py index 55f37d3..e1a470a 100644 --- a/tests/data/test_preprocess_config.py +++ b/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(): diff --git a/tests/inference/test_task.py b/tests/inference/test_task.py index ca811b1..b969c74 100644 --- a/tests/inference/test_task.py +++ b/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(): diff --git a/tests/module/test_tie_weight.py b/tests/module/test_tie_weight.py index e495b13..6a95340 100644 --- a/tests/module/test_tie_weight.py +++ b/tests/module/test_tie_weight.py @@ -1,7 +1,6 @@ import json import os -import pytest import safetensors.torch as st import torch diff --git a/tests/trainer/conftest.py b/tests/trainer/conftest.py index 39949c8..f28fe4b 100644 --- a/tests/trainer/conftest.py +++ b/tests/trainer/conftest.py @@ -1,5 +1,3 @@ -import os - import pytest import torch diff --git a/tests/trainer/test_grpo_strategy.py b/tests/trainer/test_grpo_strategy.py index 904800a..9b62434 100644 --- a/tests/trainer/test_grpo_strategy.py +++ b/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(