refactor: deduplicate preprocessing kernel and BFD packing

- Extract shared core (mask building, primary-id extraction, tensorisation, position-id generation) to astrai/preprocessing/core.py; Pipeline and TokenizeTransform both consume it, eliminating ~60% duplicated logic
- Promote BFD _plan to module-level plan_bfd(lengths, max_len) returning pure index bins; BFDPacking.apply and evaluate_ifd._pack_bins both call it, removing the second BFD implementation
- Split Pipeline._flush (49 lines) into _inject_doc_reset_position_ids + _inject_continuous_position_ids + _to_tensors; split Pipeline.run by delegating record iteration to core.iter_raw_records
- Remove dead no-op pop/塞回 in Pipeline.run (L110-111)

astrai/preprocessing/__init__.py

@@ -8,6 +8,7 @@ from astrai.preprocessing.builder import (
 from astrai.preprocessing.packing import (
     PackingStrategy,
     PackingStrategyFactory,
+    plan_bfd,
 )
 from astrai.preprocessing.pipeline import Pipeline, filter_by_length
 from astrai.preprocessing.position_id import (
@@ -35,4 +36,5 @@ __all__ = [
     "StoreWriterFactory",
     "TokenizeTransform",
     "filter_by_length",
+    "plan_bfd",
 ]

astrai/preprocessing/core.py

@@ -0,0 +1,124 @@
+"""Shared preprocessing kernel used by both :class:`Pipeline` and
+:class:`TokenizeTransform`.
+
+The two entry points previously duplicated ~60 % of their logic:
+record iteration, mask-builder invocation, primary-id extraction,
+per-key accumulation, dtype inference and position-id generation.
+This module factors out the common core as pure functions so that
+the online (``TokenizeTransform``) and offline (``Pipeline``) paths
+stay in lockstep.
+"""
+
+from itertools import chain
+from typing import Dict, Iterator, List, Optional
+
+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
+
+
+def build_preprocessing_components(config: PipelineConfig, tokenizer_path: str):
+    """Load tokenizer, mask builder and position-id strategy together.
+
+    Both ``Pipeline`` and ``TokenizeTransform`` need the same triple;
+    centralising the construction avoids drift (e.g. one path forgetting
+    to create the position-id strategy).
+    """
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
+    mask_builder = MaskBuilderFactory.create("sectioned")
+    position_strategy = PositionIdStrategyFactory.create(
+        config.output.position_ids_mode
+    )
+    return tokenizer, mask_builder, position_strategy
+
+
+def primary_ids(result: dict) -> List[int]:
+    """Return the first flat int-list value in *result*.
+
+    Used for token counting and position-id generation when the
+    primary key name is not known (DPO uses ``chosen``, GRPO uses
+    ``prompts``, SFT uses ``sequence``).
+    """
+    for val in result.values():
+        if isinstance(val, list) and val and isinstance(val[0], int):
+            return val
+    return []
+
+
+def infer_dtype(ids: List) -> torch.dtype:
+    """Float values become float32, everything else int32."""
+    if ids and isinstance(ids[0], float):
+        return torch.float32
+    return torch.int32
+
+
+def iter_raw_records(
+    records: List[dict],
+    mask_builder,
+    config: PipelineConfig,
+    tokenizer,
+) -> Iterator[dict]:
+    """Yield mask-builder output dicts for each record, skipping failures.
+
+    Drops ``domain`` from the result (callers that need it should read
+    it before calling this).  Each yielded dict maps a key
+    (``sequence``, ``chosen``, ``responses``…) to either a flat
+    ``List[int]`` or a nested ``List[List[int]]`` (GRPO responses/masks).
+    """
+    for item in records:
+        result = mask_builder.build(item, config, tokenizer)
+        if result is None:
+            continue
+        result.pop("domain", None)
+        if not primary_ids(result):
+            continue
+        yield result
+
+
+def to_per_record_tensors(
+    raw: Dict[str, list],
+) -> Dict[str, List[torch.Tensor]]:
+    """Convert an accumulated ``{key: [per-record ids]}`` dict to tensors.
+
+    Handles three shapes transparently:
+
+    - ``List[int]`` per record (``sequence``, ``chosen``…) → one tensor per record.
+    - ``List[List[int]]`` per record (GRPO ``responses``/``masks``) → one
+      ``List[Tensor]`` per record (nested), preserving the per-response
+      boundary so downstream code can index responses individually.
+    - ``List[int]`` for the whole shard (pre-packed keys) → single tensor.
+
+    The detection mirrors the previous inline logic in
+    ``Pipeline._flush`` and ``TokenizeTransform.apply``.
+    """
+    tensors: Dict[str, List[torch.Tensor]] = {}
+    for key, ids_list in raw.items():
+        if ids_list and isinstance(ids_list[0], list):
+            tensors[key] = [
+                [torch.tensor(sub, dtype=infer_dtype(sub)) for sub in ids]
+                if ids and isinstance(ids[0], list)
+                else torch.tensor(ids, dtype=infer_dtype(ids))
+                for ids in ids_list
+            ]
+        else:
+            tensors[key] = [
+                torch.tensor(list(chain.from_iterable(ids_list)), dtype=torch.int32)
+            ]
+    return tensors
+
+
+def build_position_ids(
+    sequences: List[List[int]],
+    strategy,
+) -> Optional[List[int]]:
+    """Generate position ids for *sequences* using *strategy*.
+
+    Returns ``None`` when the strategy produces no ids (e.g. ``none``
+    mode), so callers can skip attaching the key instead of storing
+    an empty list.
+    """
+    pos_ids = strategy.generate(sequences)
+    return pos_ids or None

astrai/preprocessing/packing.py

@@ -19,6 +19,43 @@ def _truncate(seq: List[int], max_len: int, mode: str) -> List[int]:
     return seq[:max_len]
 
 
+def plan_bfd(
+    sequences: List[List[int]], max_packed_len: int, truncation_mode: str = "keep_start"
+) -> List[List[int]]:
+    """Best-Fit Decreasing bin packing of *sequences* into bins.
+
+    Returns a list of bins, each bin a list of original indices into
+    *sequences*.  Bin capacities are respected on the *truncated*
+    length of each sequence (so a sequence longer than
+    *max_packed_len* counts at *max_packed_len*).
+
+    Pure index-based so callers can apply the same plan to any
+    aligned key (``loss_mask``, ``position_ids``…).
+    """
+    n = len(sequences)
+    order = sorted(range(n), key=lambda i: len(sequences[i]), reverse=True)
+    bins: List[List[int]] = []
+    bin_lengths: List[int] = []
+
+    for orig_idx in order:
+        seq_len = len(_truncate(sequences[orig_idx], max_packed_len, truncation_mode))
+        best_bin = None
+        best_remain = max_packed_len + 1
+        for i, bl in enumerate(bin_lengths):
+            remain = max_packed_len - bl
+            if seq_len <= remain < best_remain:
+                best_remain = remain
+                best_bin = i
+        if best_bin is not None:
+            bins[best_bin].append(orig_idx)
+            bin_lengths[best_bin] += seq_len
+        else:
+            bins.append([orig_idx])
+            bin_lengths.append(seq_len)
+
+    return bins
+
+
 class PackingStrategy(ABC):
     """Reorder and truncate sequences within a shard."""
 
@@ -70,7 +107,7 @@ class BFDPacking(PackingStrategy):
         sequences = keys.get("sequence", [])
         if not sequences:
             return keys
-        bins = self._plan(sequences, max_packed_len, truncation_mode)
+        bins = plan_bfd(sequences, max_packed_len, truncation_mode)
 
         packed: Dict[str, List[List[int]]] = {}
         for k, vals in keys.items():
@@ -91,35 +128,6 @@ class BFDPacking(PackingStrategy):
             result.extend(vals[i])
         return result
 
-    @staticmethod
-    def _plan(
-        sequences: List[List[int]], max_packed_len: int, truncation_mode: str
-    ) -> List[List[int]]:
-        n = len(sequences)
-        order = sorted(range(n), key=lambda i: len(sequences[i]), reverse=True)
-        bins: List[List[int]] = []
-        bin_lengths: List[int] = []
-
-        for orig_idx in order:
-            seq_len = len(
-                _truncate(sequences[orig_idx], max_packed_len, truncation_mode)
-            )
-            best_bin = None
-            best_remain = max_packed_len + 1
-            for i, bl in enumerate(bin_lengths):
-                remain = max_packed_len - bl
-                if seq_len <= remain < best_remain:
-                    best_remain = remain
-                    best_bin = i
-            if best_bin is not None:
-                bins[best_bin].append(orig_idx)
-                bin_lengths[best_bin] += seq_len
-            else:
-                bins.append([orig_idx])
-                bin_lengths.append(seq_len)
-
-        return bins
-
 
 @PackingStrategyFactory.register("bfd_split")
 class BFDSplitPacking(BFDPacking):

astrai/preprocessing/pipeline.py

@@ -4,6 +4,10 @@ Composes a :class:`BaseMaskBuilder` (selected by ``input.type``) with
 sharding and flush to ``.h5`` / ``.bin`` storage.  Packing, position-id
 generation and storage writing are each delegated to pluggable strategies,
 dispatched by configuration keys.
+
+Record iteration, mask building, primary-id extraction and per-key
+accumulation are shared with :class:`TokenizeTransform` via the
+:mod:`astrai.preprocessing.core` helpers.
 """
 
 import json
@@ -17,11 +21,13 @@ import torch
 import tqdm
 
 from astrai.config.preprocess_config import PipelineConfig
-from astrai.preprocessing.builder import MaskBuilderFactory
+from astrai.preprocessing.core import (
+    build_preprocessing_components,
+    iter_raw_records,
+    primary_ids,
+)
 from astrai.preprocessing.packing import PackingStrategyFactory
-from astrai.preprocessing.position_id import PositionIdStrategyFactory
 from astrai.preprocessing.writer import StoreWriterFactory
-from astrai.tokenize import AutoTokenizer
 
 logger = logging.getLogger(__name__)
 
@@ -64,20 +70,18 @@ class Pipeline:
         self.output_dir = output_dir
         self.tokenizer_path = tokenizer_path
 
-        self.mask_builder = MaskBuilderFactory.create("sectioned")
+        self.tokenizer, self.mask_builder, self._position_id = (
+            build_preprocessing_components(config, tokenizer_path)
+        )
         self._packer = PackingStrategyFactory.create(
             config.preprocessing.packing_strategy
         )
-        self._position_id = PositionIdStrategyFactory.create(
-            config.output.position_ids_mode
-        )
         self._writer = StoreWriterFactory.create(config.output.storage_format)
 
     def transform(self, item: dict) -> Optional[dict]:
-        return self.mask_builder.build(item, self.config, self._tokenizer)
+        return self.mask_builder.build(item, self.config, self.tokenizer)
 
     def run(self):
-        self._tokenizer = AutoTokenizer.from_pretrained(self.tokenizer_path)
         domains: dict = defaultdict(lambda: defaultdict(list))
         total_tokens = 0
         shard_idx: dict[str, int] = defaultdict(int)
@@ -102,14 +106,7 @@ class Pipeline:
                 continue
 
             domain = result.pop("domain", "__default__")
-
+            ids = primary_ids(result)
-            is_multi = bool(getattr(self.config.input, "sources", None))
-            if is_multi:
-                ids = self._primary_ids(result)
-            else:
-                ids = result.pop("sequence")
-                result["sequence"] = ids
-
             if not ids:
                 continue
 
@@ -129,15 +126,6 @@ class Pipeline:
         if total_tokens > 0:
             self._flush(domains, shard_idx)
 
-    @staticmethod
-    def _primary_ids(result: dict) -> list:
-        """Return the first list-valued entry in *result* as the primary id
-        sequence for token counting."""
-        for val in result.values():
-            if isinstance(val, list) and val and isinstance(val[0], int):
-                return val
-        return []
-
     @staticmethod
     def _align_bucket(bucket: dict, result: dict, ids: list):
         """Pad previously-accumulated keys that are missing from *result*."""
@@ -170,39 +158,12 @@ class Pipeline:
             original_sequences = keys.get("sequence", [])
             mode = self.config.output.position_ids_mode
 
-            if mode == "doc_reset" and original_sequences:
+            keys = self._inject_doc_reset_position_ids(keys, mode, original_sequences)
-                keys["position_ids"] = [list(range(len(s))) for s in original_sequences]
-
             keys = self._packer.apply(dict(keys), pp.max_packed_len, pp.truncation_mode)
-
+            tensors = self._to_tensors(keys)
-            tensors: Dict[str, List[torch.Tensor]] = {}
+            tensors = self._inject_continuous_position_ids(
-            for key, ids_list in keys.items():
+                tensors, mode, keys.get("sequence", [])
-                dt = _STR_TO_DTYPE.get(
+            )
-                    self.config.output.dtype.get(key, "int32"), torch.int32
-                )
-                # GRPO multi-response keys store List[List[int]] per record
-                # (responses/masks).  Rewards store List[float] per record.
-                # Both produce List[Tensor] (one tensor per record), but
-                # responses need inner flattening while rewards do not.
-                if ids_list and isinstance(ids_list[0], list):
-                    tensors[key] = [
-                        torch.tensor(
-                            list(chain.from_iterable(ids))
-                            if ids and isinstance(ids[0], list)
-                            else ids,
-                            dtype=dt,
-                        )
-                        for ids in ids_list
-                    ]
-                else:
-                    tensors[key] = [
-                        torch.tensor(list(chain.from_iterable(ids_list)), dtype=dt)
-                    ]
-
-            if mode == "continuous" and original_sequences:
-                pos_ids = self._position_id.generate(keys.get("sequence", []))
-                if pos_ids:
-                    tensors["position_ids"] = [torch.tensor(pos_ids, dtype=torch.int32)]
 
             self._writer.save(self.output_dir, domain, idx, tensors)
             shard_idx[domain] = idx + 1
@@ -212,3 +173,76 @@ class Pipeline:
                 f"  saved {domain}/shard_{idx:04d}  "
                 f"({tensors[first_key][0].numel():,} tokens)"
             )
+
+    def _inject_doc_reset_position_ids(
+        self,
+        keys: Dict[str, list],
+        mode: str,
+        original_sequences: List[List[int]],
+    ) -> Dict[str, list]:
+        """Attach per-document position_ids before packing (``doc_reset``).
+
+        ``doc_reset`` position ids must enter the packer so that each
+        packed bin concatenates the per-doc ranges in bin order.  The
+        per-record structure ``[range(len(s)) for s in seqs]`` is required
+        by the packer (it concatenates per-record lists per bin); the
+        ``PositionIdStrategy.generate`` flattens, so it cannot be used
+        directly here — it is only consulted for the ``continuous``
+        post-packing path.
+        """
+        if mode != "doc_reset" or not original_sequences:
+            return keys
+        keys["position_ids"] = [list(range(len(s))) for s in original_sequences]
+        return keys
+
+    def _inject_continuous_position_ids(
+        self,
+        tensors: Dict[str, List[torch.Tensor]],
+        mode: str,
+        packed_sequences: List[List[int]],
+    ) -> Dict[str, List[torch.Tensor]]:
+        """Attach a single continuous position_ids tensor after packing.
+
+        ``continuous`` mode spans the whole shard (post-packing), so it
+        cannot participate in bin packing — it is computed from the
+        packed sequences and appended directly to the tensor dict.
+        """
+        if mode != "continuous" or not packed_sequences:
+            return tensors
+        pos_ids = self._position_id.generate(packed_sequences)
+        if pos_ids:
+            tensors["position_ids"] = [torch.tensor(pos_ids, dtype=torch.int32)]
+        return tensors
+
+    def _to_tensors(self, keys: Dict[str, list]) -> Dict[str, List[torch.Tensor]]:
+        """Convert packed per-key id lists to tensors.
+
+        Honours ``config.output.dtype`` overrides per key; falls back to
+        ``int32``.  Handles three shapes (see
+        :func:`astrai.preprocessing.core.to_per_record_tensors` for the
+        equivalent online-path helper):
+        - ``List[int]`` per record → one tensor per record.
+        - ``List[List[int]]`` per record (GRPO responses/masks) → one tensor
+          per record, inner lists flattened.
+        - ``List[int]`` for the whole shard (pre-packed keys) → single tensor.
+        """
+        tensors: Dict[str, List[torch.Tensor]] = {}
+        for key, ids_list in keys.items():
+            dt = _STR_TO_DTYPE.get(
+                self.config.output.dtype.get(key, "int32"), torch.int32
+            )
+            if ids_list and isinstance(ids_list[0], list):
+                tensors[key] = [
+                    torch.tensor(
+                        list(chain.from_iterable(ids))
+                        if ids and isinstance(ids[0], list)
+                        else ids,
+                        dtype=dt,
+                    )
+                    for ids in ids_list
+                ]
+            else:
+                tensors[key] = [
+                    torch.tensor(list(chain.from_iterable(ids_list)), dtype=dt)
+                ]
+        return tensors

astrai/preprocessing/transform.py

@@ -4,6 +4,11 @@ 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.
+
+The record-processing core (mask building, primary-id extraction,
+per-key tensorisation, position-id generation) is shared with
+:class:`astrai.preprocessing.pipeline.Pipeline` via the
+:mod:`astrai.preprocessing.core` helpers.
 """
 
 import json
@@ -13,9 +18,12 @@ from typing import Dict, List
 import torch
 
 from astrai.config.preprocess_config import PipelineConfig
-from astrai.preprocessing.builder import MaskBuilderFactory
+from astrai.preprocessing.core import (
-from astrai.preprocessing.position_id import PositionIdStrategyFactory
+    build_position_ids,
-from astrai.tokenize import AutoTokenizer
+    build_preprocessing_components,
+    iter_raw_records,
+    to_per_record_tensors,
+)
 
 
 class TokenizeTransform:
@@ -33,10 +41,8 @@ class TokenizeTransform:
 
     def __init__(self, config: PipelineConfig, tokenizer_path: str):
         self.config = config
-        self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
+        self.tokenizer, self.mask_builder, self.position_strategy = (
-        self.mask_builder = MaskBuilderFactory.create("sectioned")
+            build_preprocessing_components(config, tokenizer_path)
-        self.position_strategy = PositionIdStrategyFactory.create(
-            config.output.position_ids_mode
         )
 
     @classmethod
@@ -64,40 +70,23 @@ class TokenizeTransform:
         raw: Dict[str, list] = {}
         doc_sequences: List[List[int]] = []
 
-        for item in records:
+        for result in iter_raw_records(
-            result = self.mask_builder.build(item, self.config, self.tokenizer)
+            records, self.mask_builder, self.config, self.tokenizer
-            if result is None:
+        ):
-                continue
+            primary = None
-            result.pop("domain", None)
+            for val in result.values():
-            primary_ids = self._primary_ids(result)
+                if isinstance(val, list) and val and isinstance(val[0], int):
-            if not primary_ids:
+                    primary = val
-                continue
+                    break
-            doc_sequences.append(primary_ids)
+            if primary is not None:
+                doc_sequences.append(primary)
             for key, ids in result.items():
-                if key not in raw:
+                raw.setdefault(key, []).append(ids)
-                    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)
+        tensors = to_per_record_tensors(raw)
-        if pos_ids:
-            raw["position_ids"] = [torch.tensor(pos_ids, dtype=torch.int32)]
 
-        return raw
+        pos_ids = build_position_ids(doc_sequences, self.position_strategy)
+        if pos_ids is not None:
+            tensors["position_ids"] = [torch.tensor(pos_ids, dtype=torch.int32)]
 
-    @staticmethod
+        return tensors
-    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

scripts/eval/evaluate_ifd.py

@@ -26,28 +26,22 @@ import torch.nn.functional as F
 import tqdm
 
 from astrai.model import AutoModel
+from astrai.preprocessing.packing import plan_bfd
 from astrai.tokenize import AutoTokenizer
 
 
 def _pack_bins(pairs, max_len):
-    """BFD bin packing: pack (c+r) into bins of max total length."""
+    """BFD bin packing: pack (c+r) into bins of max total length.
-    indexed = sorted(enumerate(pairs), key=lambda x: -(len(x[1][0]) + len(x[1][1])))
+
-    bins = []
+    Reuses :func:`plan_bfd` so the BFD heuristic stays single-sourced.
-    lengths = []
+    """
-    for orig_idx, (c, r) in indexed:
+    # Treat each pair as a single sequence of length len(c)+len(r) for
-        size = len(c) + len(r)
+    # planning purposes; plan_bfd works on pure lengths.
-        best_bin = -1
+    fake_sequences = [[0] * (len(c) + len(r)) for c, r in pairs]
-        for bi, rem in enumerate(lengths):
+    plan = plan_bfd(fake_sequences, max_len)
-            if rem >= size:
+    return [
-                if best_bin < 0 or rem < lengths[best_bin]:
+        [(i, pairs[i][0], pairs[i][1]) for i in bin_indices] for bin_indices in plan
-                    best_bin = bi
+    ]
-        if best_bin >= 0:
-            bins[best_bin].append((orig_idx, c, r))
-            lengths[best_bin] -= size
-        else:
-            bins.append([(orig_idx, c, r)])
-            lengths.append(max_len - size)
-    return bins
 
 
 def _resolve_sentinel_ids(tokenizer, sentinel_text):