feat: implement bfd_split packing strategy

- BFDSplitPacking splits over-length sequences into chunks before BFD
- All keys (loss_mask, position_ids, ...) split in lockstep for alignment
- No tokens lost vs bfd which truncates over-length sequences
- Tests: token preservation, chunk alignment, short unchanged, vs bfd

astrai/preprocessing/packing.py

@@ -119,3 +119,50 @@ class BFDPacking(PackingStrategy):
                 bin_lengths.append(seq_len)
 
         return bins
+
+
+@PackingStrategyFactory.register("bfd_split")
+class BFDSplitPacking(BFDPacking):
+    """BFD packing with over-length sequences split into chunks.
+
+    Sequences longer than *max_packed_len* are split into consecutive
+    chunks of at most *max_packed_len* tokens instead of being
+    truncated.  Each chunk becomes an independent sequence that enters
+    BFD planning.  All keys (``loss_mask``, ``position_ids``, …) are
+    split in lockstep so per-token alignment is preserved.
+
+    Note: because each chunk is treated as a separate document, the
+    second chunk of a split sequence loses the preceding context.
+    """
+
+    def apply(
+        self,
+        keys: Dict[str, List[List[int]]],
+        max_packed_len: int,
+        truncation_mode: str,
+    ) -> Dict[str, List[List[int]]]:
+        sequences = keys.get("sequence", [])
+        if not sequences:
+            return keys
+        if max_packed_len <= 0:
+            return super().apply(keys, max_packed_len, truncation_mode)
+
+        split_keys = self._split_all(keys, max_packed_len)
+        return super().apply(split_keys, max_packed_len, truncation_mode)
+
+    @staticmethod
+    def _split_all(
+        keys: Dict[str, List[List[int]]], max_packed_len: int
+    ) -> Dict[str, List[List[int]]]:
+        """Split every sequence exceeding *max_packed_len* into chunks,
+        applying the same chunk boundaries to all keys."""
+        sequences = keys["sequence"]
+        chunk_bounds = [list(range(0, len(s), max_packed_len)) for s in sequences]
+        result: Dict[str, List[List[int]]] = {}
+        for key, vals in keys.items():
+            split_vals: List[List[int]] = []
+            for val, starts in zip(vals, chunk_bounds):
+                for start in starts:
+                    split_vals.append(val[start : start + max_packed_len])
+            result[key] = split_vals
+        return result

tests/data/test_preprocess_pipeline.py

@@ -7,6 +7,7 @@ from astrai.config.preprocess_config import (
     PipelineConfig,
     ProcessingConfig,
 )
+from astrai.preprocessing.packing import PackingStrategyFactory
 from astrai.preprocessing.pipeline import Pipeline, filter_by_length
 from tests.data.conftest import (
     _CHAT_SECTIONS,
@@ -262,3 +263,69 @@ def test_grpo_pipeline(temp_dir, tokenizer_dir):
     assert "masks" in meta
     assert "rewards" in meta
     assert "sequence" not in meta
+
+
+# ---------------------------------------------------------------------------
+# BFD split packing
+# ---------------------------------------------------------------------------
+
+_TRU = "keep_start"
+
+
+def _total_tokens(keys, key="sequence"):
+    return sum(len(s) for s in keys[key])
+
+
+def test_bfd_split_preserves_all_tokens():
+    """No tokens are lost — split chunks are kept, not truncated away."""
+    packer = PackingStrategyFactory.create("bfd_split")
+    max_len = 10
+    keys = {
+        "sequence": [list(range(25)), list(range(3))],
+        "loss_mask": [[1] * 25, [1] * 3],
+    }
+    result = packer.apply(keys, max_len, _TRU)
+
+    assert _total_tokens(result) == 28
+    for seq in result["sequence"]:
+        assert len(seq) <= max_len
+
+
+def test_bfd_split_chunk_alignment():
+    """loss_mask chunks must align with sequence chunks."""
+    packer = PackingStrategyFactory.create("bfd_split")
+    max_len = 10
+    keys = {
+        "sequence": [list(range(25))],
+        "loss_mask": [[0] * 5 + [1] * 20],
+    }
+    result = packer.apply(keys, max_len, _TRU)
+
+    for seq, mask in zip(result["sequence"], result["loss_mask"]):
+        assert len(seq) == len(mask)
+
+
+def test_bfd_split_short_unchanged():
+    """Sequences under max_packed_len should not be split."""
+    packer = PackingStrategyFactory.create("bfd_split")
+    max_len = 10
+    keys = {"sequence": [list(range(5))], "loss_mask": [[1] * 5]}
+    result = packer.apply(keys, max_len, _TRU)
+
+    assert _total_tokens(result) == 5
+    assert len(result["sequence"]) >= 1
+
+
+def test_bfd_split_vs_bfd():
+    """bfd loses tokens from over-length sequences; bfd_split does not."""
+    max_len = 10
+    keys = {
+        "sequence": [list(range(25)), list(range(8))],
+        "loss_mask": [[1] * 25, [1] * 8],
+    }
+
+    bfd = PackingStrategyFactory.create("bfd").apply(keys, max_len, _TRU)
+    split = PackingStrategyFactory.create("bfd_split").apply(keys, max_len, _TRU)
+
+    assert _total_tokens(bfd) < 33
+    assert _total_tokens(split) == 33