feat: add JSONL dataset store with on-the-fly tokenization

- Add JsonlStore registered under "jsonl" in astrai/dataset/storage.py
- Reuse PipelineConfig schema for JSONL dataset configuration
- Update detect_format to recognize JSONL directories and files
- Move save_h5/load_h5/save_bin/load_bin to astrai/serialization
- Split astrai/serialization.py into checkpoint/dataset submodules
- Add tests for JSONL detection, seq/SFT stores, and config roundtrip

astrai/dataset/__init__.py

@@ -5,10 +5,13 @@ from astrai.dataset.dataset import (
 from astrai.dataset.sampler import ResumableDistributedSampler
 from astrai.dataset.storage import (
     H5Store,
+    JsonlStore,
     MmapStore,
     Store,
     StoreFactory,
     detect_format,
+)
+from astrai.serialization import (
     load_bin,
     load_h5,
     save_bin,
@@ -22,6 +25,7 @@ __all__ = [
     "StoreFactory",
     "H5Store",
     "MmapStore",
+    "JsonlStore",
     "detect_format",
     "save_h5",
     "load_h5",

astrai/dataset/dataset.py

@@ -48,24 +48,26 @@ class BaseDataset(Dataset, ABC):
                 f"Missing: {missing}"
             )
 
-    def load(self, load_path: str, storage_type: Optional[str] = None):
+    def load(self, load_path: str, storage_type: Optional[str] = None, **kwargs):
         """Load dataset from the given path.
 
         Auto-detects the storage format if not specified.
 
         Args:
             load_path: Path to the data directory or file
-            storage_type: Force a specific storage type ("h5", "bin"),
+            storage_type: Force a specific storage type ("h5", "bin", "jsonl"),
                           or None for auto-detection
+            **kwargs: Extra arguments forwarded to the store constructor and
+                      to ``store.load()``.
 
         Raises:
             KeyError: If the loaded storage is missing required keys.
         """
         if storage_type is None:
             storage_type = detect_format(load_path)
-        self.storage = StoreFactory.create(storage_type)
+        self.storage = StoreFactory.create(storage_type, **kwargs)
         self._load_path = load_path
-        self.storage.load(load_path)
+        self.storage.load(load_path, **kwargs)
         self._validate_keys()
 
     @property
@@ -144,6 +146,7 @@ class DatasetFactory(BaseFactory["BaseDataset"]):
         window_size: int,
         stride: Optional[int] = None,
         storage_type: Optional[str] = None,
+        **kwargs,
     ) -> "BaseDataset":
         """Create and load a dataset in one step.
 
@@ -152,7 +155,8 @@ class DatasetFactory(BaseFactory["BaseDataset"]):
             load_path: Path to the data file
             window_size: Window size for data sampling
             stride: Stride between consecutive samples (default: same as window_size)
-            storage_type: Storage type ("h5", "bin") or None for auto-detection
+            storage_type: Storage type ("h5", "bin", "jsonl") or None for auto-detection
+            **kwargs: Extra arguments forwarded to ``dataset.load()``.
 
         Returns:
             Loaded dataset instance
@@ -161,7 +165,7 @@ class DatasetFactory(BaseFactory["BaseDataset"]):
             stride = window_size
 
         dataset = cls.create(train_type, window_size, stride)
-        dataset.load(load_path, storage_type=storage_type)
+        dataset.load(load_path, storage_type=storage_type, **kwargs)
 
         return dataset
 

astrai/dataset/storage.py

@@ -20,79 +20,25 @@ Key properties:
 import bisect
 import glob
 import json
-import os
+import logging
 from abc import ABC, abstractmethod
 from pathlib import Path
 from typing import Dict, List, Union
 
-import h5py
-import numpy as np
 import torch
 from torch import Tensor
 
+from astrai.config.preprocess_config import PipelineConfig
 from astrai.factory import BaseFactory
+from astrai.preprocessing.builder import MaskBuilderFactory
+from astrai.preprocessing.position_id import PositionIdStrategyFactory
+from astrai.serialization import (
+    load_bin,
+    load_h5,
+)
+from astrai.tokenize import AutoTokenizer
 
-
+logger = logging.getLogger(__name__)
-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)
-    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]]):
-    os.makedirs(file_path, exist_ok=True)
-    meta = {}
-    for key, tensors in tensor_group.items():
-        cat = torch.cat(tensors, dim=0)
-        meta[key] = {"shape": list(cat.shape), "dtype": str(cat.dtype).split(".")[-1]}
-        np.asarray(cat.cpu().numpy()).tofile(os.path.join(file_path, f"{key}.bin"))
-    with open(os.path.join(file_path, "meta.json"), "w") as f:
-        json.dump(meta, f)
-
-
-def load_bin(file_path: str) -> Dict[str, List[Tensor]]:
-    with open(os.path.join(file_path, "meta.json"), "r") as f:
-        meta = json.load(f)
-    segments: Dict[str, List[Tensor]] = {}
-    for key, info in meta.items():
-        arr = np.memmap(
-            os.path.join(file_path, f"{key}.bin"),
-            dtype=info["dtype"],
-            mode="r+",
-            shape=tuple(info["shape"]),
-        )
-        segments[key] = [torch.from_numpy(arr)]
-    return segments
 
 
 def detect_format(load_path: str) -> str:
@@ -102,7 +48,7 @@ def detect_format(load_path: str) -> str:
         load_path: Directory or file path
 
     Returns:
-        Format string ("h5" or "bin")
+        Format string ("h5", "bin", or "jsonl")
 
     Raises:
         FileNotFoundError: If no supported data files are found
@@ -112,6 +58,8 @@ def detect_format(load_path: str) -> str:
         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 = [
@@ -128,6 +76,11 @@ def detect_format(load_path: str) -> str:
         ) > 0
         if has_meta:
             return "bin"
+    jsonl_files = [
+        Path(p) for p in glob.glob(str(root / "**" / "*.jsonl"), recursive=True)
+    ]
+    if jsonl_files:
+        return "jsonl"
     raise FileNotFoundError(f"No supported data files found at {load_path}")
 
 
@@ -264,3 +217,96 @@ class MmapStore(Store):
         self._normalize(all_raw)
         for tensors in self._data.values():
             self._mmap_refs.extend(tensors)
+
+
+@StoreFactory.register("jsonl")
+class JsonlStore(Store):
+    """On-the-fly tokenization store for raw JSONL files.
+
+    A JSONL dataset directory contains ``*.jsonl`` files plus a
+    ``dataset_config.json`` file that follows the same schema as
+    :class:`PipelineConfig` with an additional ``tokenizer_path`` field.
+    Records are tokenized when the store is loaded and concatenated into
+    segmented tensors matching the key layout expected by the dataset
+    classes (``sequence``, ``loss_mask``, ``position_ids``, ...).
+    """
+
+    CONFIG_NAME = "dataset_config.json"
+
+    def load(self, path: str):
+        root = Path(path)
+        config_path = root / self.CONFIG_NAME
+        if not config_path.exists():
+            raise FileNotFoundError(
+                f"JSONL dataset config not found: {config_path}. "
+                f"Expected {self.CONFIG_NAME} alongside *.jsonl files."
+            )
+
+        with open(config_path, "r", encoding="utf-8") as f:
+            raw_config = json.load(f)
+
+        tokenizer_path = raw_config.pop("tokenizer_path", None)
+        if tokenizer_path is None:
+            raise ValueError(
+                f"JSONL dataset config must specify 'tokenizer_path': {config_path}"
+            )
+
+        self.config = PipelineConfig.from_dict(raw_config)
+        tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
+        mask_builder = MaskBuilderFactory.create("sectioned")
+        position_strategy = PositionIdStrategyFactory.create(
+            self.config.output.position_ids_mode
+        )
+
+        raw: Dict[str, List[Tensor]] = {}
+        doc_sequences: List[List[int]] = []
+
+        for jsonl_path in sorted(root.glob("*.jsonl")):
+            with open(jsonl_path, "r", encoding="utf-8") as f:
+                for line in f:
+                    line = line.strip()
+                    if not line:
+                        continue
+                    try:
+                        item = json.loads(line)
+                    except json.JSONDecodeError:
+                        logger.warning(
+                            "Failed to parse JSON line in %s, skipping", jsonl_path
+                        )
+                        continue
+
+                    result = mask_builder.build(item, self.config, tokenizer)
+                    if result is None:
+                        continue
+
+                    result.pop("domain", None)
+                    primary_ids = self._primary_ids(result)
+                    if not primary_ids:
+                        continue
+
+                    doc_sequences.append(primary_ids)
+                    for key, ids in result.items():
+                        if key not in raw:
+                            raw[key] = []
+                        raw[key].append(torch.tensor(ids, dtype=self._infer_dtype(ids)))
+
+        pos_ids = position_strategy.generate(doc_sequences)
+        if pos_ids:
+            raw["position_ids"] = [torch.tensor(pos_ids, dtype=torch.int32)]
+
+        self._normalize(raw)
+
+    @staticmethod
+    def _primary_ids(result: dict) -> List[int]:
+        """Return the first integer list in *result* as the primary id sequence."""
+        for val in result.values():
+            if isinstance(val, list) and val and isinstance(val[0], int):
+                return val
+        return []
+
+    @staticmethod
+    def _infer_dtype(ids: List) -> torch.dtype:
+        """Infer tensor dtype from the first element of a token/value list."""
+        if ids and isinstance(ids[0], float):
+            return torch.float32
+        return torch.int32

astrai/preprocessing/writer.py

@@ -14,8 +14,8 @@ from typing import Dict, List
 
 import torch
 
-from astrai.dataset.storage import save_bin, save_h5
 from astrai.factory import BaseFactory
+from astrai.serialization import save_bin, save_h5
 
 logger = logging.getLogger(__name__)
 

astrai/serialization/__init__.py

@@ -0,0 +1,43 @@
+"""Serialization utilities for models and datasets.
+
+This package re-exports checkpoint helpers and dataset storage helpers so
+that existing imports from ``astrai.serialization`` continue to work.
+"""
+
+from astrai.serialization.checkpoint import (
+    Checkpoint,
+    load_json,
+    load_model_config,
+    load_model_weights,
+    load_safetensors,
+    load_state_dict,
+    load_torch,
+    save_json,
+    save_model,
+    save_safetensors,
+    save_torch,
+)
+from astrai.serialization.dataset import (
+    load_bin,
+    load_h5,
+    save_bin,
+    save_h5,
+)
+
+__all__ = [
+    "Checkpoint",
+    "load_json",
+    "load_model_config",
+    "load_model_weights",
+    "load_safetensors",
+    "load_state_dict",
+    "load_torch",
+    "save_json",
+    "save_model",
+    "save_safetensors",
+    "save_torch",
+    "load_bin",
+    "load_h5",
+    "save_bin",
+    "save_h5",
+]

astrai/serialization/checkpoint.py

@@ -1,5 +1,8 @@
+"""Model checkpoint serialization helpers."""
+
 import io
 import json
+import os
 import time
 from dataclasses import dataclass, field
 from pathlib import Path

astrai/serialization/dataset.py

@@ -0,0 +1,73 @@
+"""Dataset storage serialization helpers (HDF5 / memory-mapped binary)."""
+
+import json
+import os
+from pathlib import Path
+from typing import Dict, List
+
+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)
+    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]]):
+    os.makedirs(file_path, exist_ok=True)
+    meta = {}
+    for key, tensors in tensor_group.items():
+        cat = torch.cat(tensors, dim=0)
+        meta[key] = {"shape": list(cat.shape), "dtype": str(cat.dtype).split(".")[-1]}
+        np.asarray(cat.cpu().numpy()).tofile(os.path.join(file_path, f"{key}.bin"))
+    with open(os.path.join(file_path, "meta.json"), "w") as f:
+        json.dump(meta, f)
+
+
+def load_bin(file_path: str) -> Dict[str, List[Tensor]]:
+    with open(os.path.join(file_path, "meta.json"), "r") as f:
+        meta = json.load(f)
+    segments: Dict[str, List[Tensor]] = {}
+    for key, info in meta.items():
+        arr = np.memmap(
+            os.path.join(file_path, f"{key}.bin"),
+            dtype=info["dtype"],
+            mode="r+",
+            shape=tuple(info["shape"]),
+        )
+        segments[key] = [torch.from_numpy(arr)]
+    return segments

tests/data/test_dataset.py

@@ -1,14 +1,18 @@
+import json
 import os
 
 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.storage import (
     H5Store,
     StoreFactory,
     detect_format,
+)
+from astrai.serialization import (
     load_bin,
     save_bin,
     save_h5,
@@ -19,6 +23,39 @@ def _rand_seq(length, vocab=1000):
     return torch.randint(0, vocab, (length,), dtype=torch.int64)
 
 
+def _save_test_tokenizer(test_dir, tokenizer):
+    tokenizer_path = os.path.join(test_dir, "tokenizer")
+    os.makedirs(tokenizer_path, exist_ok=True)
+    tokenizer.save_pretrained(tokenizer_path)
+    return tokenizer_path
+
+
+def _write_jsonl_dataset(test_dir, tokenizer_path, records, config_overrides=None):
+    data_dir = os.path.join(test_dir, "jsonl_data")
+    os.makedirs(data_dir, exist_ok=True)
+
+    with open(os.path.join(data_dir, "data.jsonl"), "w", encoding="utf-8") as f:
+        for record in records:
+            f.write(json.dumps(record, ensure_ascii=False) + "\n")
+
+    config = {
+        "tokenizer_path": tokenizer_path,
+        "version": 1,
+        "input": {"sections": [{"field": "text", "action": "train"}]},
+        "preprocessing": {"max_seq_len": 128},
+        "output": {"position_ids_mode": "continuous"},
+    }
+    if config_overrides:
+        config.update(config_overrides)
+
+    with open(
+        os.path.join(data_dir, "dataset_config.json"), "w", encoding="utf-8"
+    ) as f:
+        json.dump(config, f, ensure_ascii=False, indent=2)
+
+    return data_dir
+
+
 def _make_seq_dataset(
     test_dir, name="data", seq_length=200, train_type="seq", data=None, **load_kwargs
 ):
@@ -372,3 +409,106 @@ def test_dataset_load_explicit_storage_type(base_test_env):
     dataset = _make_seq_dataset(test_dir, "explicit", storage_type="h5")
     assert len(dataset) > 0
     assert dataset.count == 200
+
+
+def test_detect_format_jsonl_dir(base_test_env):
+    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 baz"}],
+    )
+    assert detect_format(data_dir) == "jsonl"
+
+
+def test_jsonl_store_seq(base_test_env):
+    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 baz qux"}],
+        config_overrides={"preprocessing": {"max_seq_len": 128, "min_chars": 0}},
+    )
+
+    store = StoreFactory.create("jsonl")
+    store.load(data_dir)
+    assert len(store) > 0
+    assert "sequence" in store.keys
+
+    dataset = DatasetFactory.load("seq", data_dir, window_size=8)
+    assert len(dataset) > 0
+    item = dataset[0]
+    assert "input_ids" in item
+    assert "target_ids" in item
+    assert item["input_ids"].dtype == torch.long
+
+
+def test_jsonl_store_sft(base_test_env):
+    test_dir = base_test_env["test_dir"]
+    tokenizer = base_test_env["tokenizer"]
+    tokenizer.set_chat_template(
+        "{% for message in messages %}{{ message['role'] }}:{{ message['content'] }}\n{% endfor %}"
+    )
+    tokenizer_path = _save_test_tokenizer(test_dir, tokenizer)
+    data_dir = _write_jsonl_dataset(
+        test_dir,
+        tokenizer_path,
+        [
+            {
+                "messages": [
+                    {"role": "system", "content": "sys"},
+                    {"role": "user", "content": "hi"},
+                    {"role": "assistant", "content": "hello"},
+                ]
+            }
+        ],
+        config_overrides={
+            "input": {
+                "sections": [{"field": "messages", "action": "$role", "template": True}]
+            },
+            "mask": {"system": "mask", "user": "mask", "assistant": "train"},
+            "mask_default": "mask",
+        },
+    )
+
+    store = StoreFactory.create("jsonl")
+    store.load(data_dir)
+    assert "sequence" in store.keys
+    assert "loss_mask" in store.keys
+    assert "position_ids" in store.keys
+
+    dataset = DatasetFactory.load("sft", data_dir, window_size=8)
+    item = dataset[0]
+    assert "input_ids" in item
+    assert "target_ids" in item
+    assert "loss_mask" in item
+    assert "position_ids" in item
+    assert item["loss_mask"].dtype == torch.bool
+
+
+def test_jsonl_store_pipeline_config_roundtrip(base_test_env):
+    test_dir = base_test_env["test_dir"]
+    config_path = os.path.join(test_dir, "dataset_config.json")
+    with open(config_path, "w", encoding="utf-8") as f:
+        json.dump(
+            {
+                "tokenizer_path": os.path.join(test_dir, "tokenizer"),
+                "version": 1,
+                "input": {"sections": [{"field": "text", "action": "train"}]},
+                "mask": {"assistant": "train"},
+                "preprocessing": {"max_seq_len": 64},
+                "output": {"position_ids_mode": "doc_reset"},
+            },
+            f,
+            ensure_ascii=False,
+            indent=2,
+        )
+
+    with open(config_path, "r", encoding="utf-8") as f:
+        raw = json.load(f)
+    raw.pop("tokenizer_path")
+    config = PipelineConfig.from_dict(raw)
+    assert config.output.position_ids_mode == "doc_reset"
+    assert config.preprocessing.max_seq_len == 64