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
This commit is contained in:
2026-07-29 12:50:27 +08:00
parent c2b04d8458
commit 115192c67c
17 changed files with 91 additions and 209 deletions
+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