diff --git a/astrai/preprocessing/packing.py b/astrai/preprocessing/packing.py index 1d300ee..05e2621 100644 --- a/astrai/preprocessing/packing.py +++ b/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 diff --git a/tests/data/test_preprocess_pipeline.py b/tests/data/test_preprocess_pipeline.py index 250de4b..90b7397 100644 --- a/tests/data/test_preprocess_pipeline.py +++ b/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