refactor: decouple tokenizer from Store into Transform layer

- Extract tokenization/mask/position logic from JsonlStore into TokenizeTransform
- JsonlStore now pure reader: reads JSON records, delegates to transform
- Store no longer imports tokenizer or preprocessing components
- Replace per_record param with segments_are_records class attribute
- Store subclasses declare segment semantics as format-level property

astrai/dataset/storage.py

@@ -28,16 +28,13 @@ from typing import Dict, List, Optional, Union
 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.transform import TokenizeTransform
-from astrai.preprocessing.position_id import PositionIdStrategyFactory
 from astrai.serialization import (
     load_bin,
     load_bin_offsets,
     load_h5,
 )
-from astrai.tokenize import AutoTokenizer
 
 logger = logging.getLogger(__name__)
 
@@ -104,10 +101,13 @@ class Store(ABC):
         count.  Backed by either per-record segment lists (H5/JSONL) or
         a single concatenated segment plus per-record offsets (bin).
 
-    Subclasses fill ``self._data`` and ``self._cum`` during ``load()``
+    Subclasses declare ``segments_are_records`` to indicate whether their
-    via ``_normalize()``.
+    segments are inherently per-record (H5/JSONL) or opaque shards (bin).
+    This is a format-level property, not a per-call decision.
     """
 
+    segments_are_records: bool = False
+
     def __init__(self):
         self._data: Dict[str, List[Tensor]] = {}
         self._cum: Dict[str, List[int]] = {}
@@ -208,7 +208,6 @@ class Store(ABC):
         self,
         raw: Dict[str, list],
         offsets: Optional[Dict[str, List[int]]] = None,
-        per_record: bool = False,
     ):
         """Register segments and pre-compute indices for both access modes.
 
@@ -217,8 +216,9 @@ class Store(ABC):
 
         Record mode: if *offsets* is provided (bin layout), ``_offsets[key]``
         stores cumulative per-record offsets into the single concatenated
-        segment.  Otherwise (h5/jsonl layout), ``_data[key]`` is already a
+        segment.  Otherwise, when ``segments_are_records`` is True
-        per-record list and ``fetch_record`` indexes it directly.
+        (H5/JSONL), ``_data[key]`` is a per-record list and
+        ``fetch_record`` indexes it directly.
 
         Nested keys (GRPO ``responses``/``masks`` as ``List[List[Tensor]]``)
         are stored as-is and excluded from both cumulative bookkeepings —
@@ -265,10 +265,7 @@ class Store(ABC):
         if valid_offsets:
             record_counts = [len(v) - 1 for v in valid_offsets.values()]
             self._num_records = min(record_counts) if record_counts else 0
-        elif per_record:
+        elif self.segments_are_records:
-            # H5/JSONL layout: _data[key] is a per-record list where each
-            # segment is one record.  Even a single segment counts as one
-            # record.
             per_record_counts = []
             for key, tensors in self._data.items():
                 if not tensors or isinstance(tensors[0], list):
@@ -276,8 +273,6 @@ class Store(ABC):
                 per_record_counts.append(len(tensors))
             self._num_records = min(per_record_counts) if per_record_counts else 0
         else:
-            # bin layout without offsets: _data[key] is [concatenated_stream].
-            # Cannot determine record boundaries — stream mode only.
             self._num_records = 0
 
 
@@ -301,8 +296,10 @@ class H5Store(Store):
     Stream mode concatenates across records via ``_cum``.
     """
 
+    segments_are_records = True
+
     def load(self, path: str):
-        self._normalize(load_h5(path), per_record=True)
+        self._normalize(load_h5(path))
 
 
 @StoreFactory.register("bin")
@@ -353,62 +350,46 @@ class MmapStore(Store):
 
 @StoreFactory.register("jsonl")
 class JsonlStore(Store):
-    """On-the-fly tokenization store for raw JSONL files.
+    """JSONL reader with pluggable tokenization transform.
 
     A JSONL dataset directory contains ``*.jsonl`` files plus a
-    ``dataset_config.json`` file that follows the same schema as
+    ``dataset_config.json`` describing the tokenization pipeline.
-    :class:`PipelineConfig` with an additional ``tokenizer_path`` field.
+
-    Records are tokenized when the store is loaded and concatenated into
+    Responsibilities are split across two layers:
-    segmented tensors matching the key layout expected by the dataset
+
-    classes (``sequence``, ``loss_mask``, ``position_ids``, ...).
+    - **Reader** (this class): reads raw JSON records from disk.
+    - **Transform** (:class:`~astrai.preprocessing.transform.TokenizeTransform`):
+      tokenizes records into per-key tensors.  When not supplied explicitly,
+      a default transform is built from ``dataset_config.json`` so that
+      existing callers keep working without changes.
+
+    The Store itself never imports the tokenizer — the dependency lives
+    in the Transform layer.
     """
 
     CONFIG_NAME = "dataset_config.json"
+    segments_are_records = True
 
-    def load(self, path: str):
+    def load(self, path: str, transform=None, **kwargs):
         root = Path(path)
-        config_path = root / self.CONFIG_NAME
+        records = self._read_records(root)
-        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:
+        if transform is None:
-            raw_config = json.load(f)
+            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, "
+                    f"or pass an explicit transform."
+                )
+            transform = TokenizeTransform.from_config_file(str(config_path))
 
-        tokenizer_path = raw_config.pop("tokenizer_path", None) or str(root)
+        raw = transform.apply(records)
-        self.config = PipelineConfig.from_dict(raw_config)
+        self._normalize(raw)
-        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]] = []
-
-        def _process_item(item: dict) -> None:
-            nonlocal raw, doc_sequences
-            result = mask_builder.build(item, self.config, tokenizer)
-            if result is None:
-                return
-            result.pop("domain", None)
-            primary_ids = self._primary_ids(result)
-            if not primary_ids:
-                return
-            doc_sequences.append(primary_ids)
-            for key, ids in result.items():
-                if key not in raw:
-                    raw[key] = []
-                if ids and isinstance(ids[0], list):
-                    # GRPO multi-response: List[List[int]] → List[Tensor]
-                    raw[key].append(
-                        [torch.tensor(sub, dtype=self._infer_dtype(sub)) for sub in ids]
-                    )
-                else:
-                    raw[key].append(torch.tensor(ids, dtype=self._infer_dtype(ids)))
 
+    @staticmethod
+    def _read_records(root: Path) -> List[dict]:
+        records: List[dict] = []
         for jsonl_path in sorted(root.glob("*.jsonl")):
             with open(jsonl_path, "r", encoding="utf-8") as f:
                 for line in f:
@@ -416,16 +397,13 @@ class JsonlStore(Store):
                     if not line:
                         continue
                     try:
-                        item = json.loads(line)
+                        records.append(json.loads(line))
                     except json.JSONDecodeError:
                         logger.warning(
                             "Failed to parse JSON line in %s, skipping", jsonl_path
                         )
-                        continue
-                    _process_item(item)
-
         for json_path in sorted(root.glob("*.json")):
-            if json_path.name == self.CONFIG_NAME:
+            if json_path.name == JsonlStore.CONFIG_NAME:
                 continue
             with open(json_path, "r", encoding="utf-8") as f:
                 try:
@@ -434,28 +412,7 @@ class JsonlStore(Store):
                     logger.warning("Failed to parse JSON file %s, skipping", json_path)
                     continue
             if isinstance(data, list):
-                for item in data:
+                records.extend(data)
-                    _process_item(item)
             elif isinstance(data, dict):
-                _process_item(data)
+                records.append(data)
-
+        return records
-        pos_ids = position_strategy.generate(doc_sequences)
-        if pos_ids:
-            raw["position_ids"] = [torch.tensor(pos_ids, dtype=torch.int32)]
-
-        self._normalize(raw, per_record=True)
-
-    @staticmethod
-    def _primary_ids(result: dict) -> List[int]:
-        """Return the first flat 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/__init__.py

@@ -14,6 +14,7 @@ from astrai.preprocessing.position_id import (
     PositionIdStrategy,
     PositionIdStrategyFactory,
 )
+from astrai.preprocessing.transform import TokenizeTransform
 from astrai.preprocessing.writer import (
     StoreWriter,
     StoreWriterFactory,
@@ -32,5 +33,6 @@ __all__ = [
     "SingleOutputMaskBuilder",
     "StoreWriter",
     "StoreWriterFactory",
+    "TokenizeTransform",
     "filter_by_length",
 ]

astrai/preprocessing/transform.py

@@ -0,0 +1,103 @@
+"""Tokenization transform for JSONL record streams.
+
+Bridges the Reader layer (``JsonlStore`` reads raw JSON records) and the
+Dataset layer (expects per-record tensors).  Holds the tokenizer,
+mask-builder and position-id strategy together so that I/O code stays
+free of model dependencies.
+"""
+
+import json
+from pathlib import Path
+from typing import Dict, List
+
+import torch
+
+from astrai.config.preprocess_config import PipelineConfig
+from astrai.preprocessing.builder import MaskBuilderFactory
+from astrai.preprocessing.position_id import PositionIdStrategyFactory
+from astrai.tokenize import AutoTokenizer
+
+
+class TokenizeTransform:
+    """Tokenize raw JSONL record dicts into per-key tensor lists.
+
+    Owns the three preprocessing concerns that were previously inlined in
+    ``JsonlStore``: tokenization, loss-mask construction and position-id
+    generation.  Constructing it loads the tokenizer, so it is intentionally
+    cheap to pass around once built.
+
+    Args:
+        config: Pipeline config describing sections / masks / position mode.
+        tokenizer_path: Path passed to ``AutoTokenizer.from_pretrained``.
+    """
+
+    def __init__(self, config: PipelineConfig, tokenizer_path: str):
+        self.config = config
+        self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
+        self.mask_builder = MaskBuilderFactory.create("sectioned")
+        self.position_strategy = PositionIdStrategyFactory.create(
+            config.output.position_ids_mode
+        )
+
+    @classmethod
+    def from_config_file(cls, config_path: str) -> "TokenizeTransform":
+        """Build from a ``dataset_config.json`` file path.
+
+        The config file follows :class:`PipelineConfig` schema with an
+        extra ``tokenizer_path`` field.  When omitted, the config's
+        parent directory is used as the tokenizer path.
+        """
+        root = Path(config_path).parent
+        with open(config_path, "r", encoding="utf-8") as f:
+            raw_config = json.load(f)
+        tokenizer_path = raw_config.pop("tokenizer_path", None) or str(root)
+        config = PipelineConfig.from_dict(raw_config)
+        return cls(config, tokenizer_path)
+
+    def apply(self, records: List[dict]) -> Dict[str, list]:
+        """Tokenize a list of raw record dicts.
+
+        Returns a dict mapping key (``sequence``, ``chosen``, ``responses``,
+        …) to a list of per-record tensors (or nested tensor lists for
+        multi-response keys such as GRPO ``responses``).
+        """
+        raw: Dict[str, list] = {}
+        doc_sequences: List[List[int]] = []
+
+        for item in records:
+            result = self.mask_builder.build(item, self.config, self.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] = []
+                if ids and isinstance(ids[0], list):
+                    raw[key].append(
+                        [torch.tensor(sub, dtype=self._infer_dtype(sub)) for sub in ids]
+                    )
+                else:
+                    raw[key].append(torch.tensor(ids, dtype=self._infer_dtype(ids)))
+
+        pos_ids = self.position_strategy.generate(doc_sequences)
+        if pos_ids:
+            raw["position_ids"] = [torch.tensor(pos_ids, dtype=torch.int32)]
+
+        return raw
+
+    @staticmethod
+    def _primary_ids(result: dict) -> List[int]:
+        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:
+        if ids and isinstance(ids[0], float):
+            return torch.float32
+        return torch.int32