refactor: split mask builder by single/multi output

- Extract SingleOutputMaskBuilder for SFT and pretrain configs
- Extract MultiOutputMaskBuilder for DPO and GRPO configs
- Keep SectionedMaskBuilder as backward-compatible facade
- Register "single" and "multi" names in MaskBuilderFactory
- Add parity and rejection tests for concrete builders

astrai/preprocessing/__init__.py

@@ -1,7 +1,9 @@
 from astrai.preprocessing.builder import (
     BaseMaskBuilder,
     MaskBuilderFactory,
+    MultiOutputMaskBuilder,
     SectionedMaskBuilder,
+    SingleOutputMaskBuilder,
 )
 from astrai.preprocessing.packing import (
     PackingStrategy,
@@ -20,12 +22,14 @@ from astrai.preprocessing.writer import (
 __all__ = [
     "BaseMaskBuilder",
     "MaskBuilderFactory",
+    "MultiOutputMaskBuilder",
     "PackingStrategy",
     "PackingStrategyFactory",
     "Pipeline",
     "PositionIdStrategy",
     "PositionIdStrategyFactory",
     "SectionedMaskBuilder",
+    "SingleOutputMaskBuilder",
     "StoreWriter",
     "StoreWriterFactory",
     "filter_by_length",

astrai/preprocessing/builder.py

@@ -1,8 +1,10 @@
 """Mask building for preprocessing pipeline.
 
 :class:`SectionRenderer` converts section specs into token ids and loss
-masks (template / text / value extraction).  :class:`SectionedMaskBuilder`
+masks (template / text / value extraction).  :class:`SingleOutputMaskBuilder`
-orchestrates single-output / multi-output (DPO / GRPO) assembly.
+handles single-output (SFT / pretrain), :class:`MultiOutputMaskBuilder`
+handles multi-output (DPO / GRPO), and :class:`SectionedMaskBuilder`
+orchestrates both modes as a façade.
 """
 
 from abc import ABC, abstractmethod
@@ -212,42 +214,17 @@ class MaskBuilderFactory(BaseFactory["BaseMaskBuilder"]):
     pass
 
 
-@MaskBuilderFactory.register("sectioned")
+@MaskBuilderFactory.register("single")
-class SectionedMaskBuilder(BaseMaskBuilder):
+class SingleOutputMaskBuilder(BaseMaskBuilder):
-    """Config-driven builder supporting single and multi-output modes.
+    """Build a single output sequence with optional loss mask.
 
-    Single-output::
+    Expects ``config.input.sections`` (list of section specs).
-
-        {"input": {"sections": [
-            {"field": "messages", "action": "$role", "template": true}
-        ]}}
-        → {"sequence": [...], "loss_mask": [...], "domain": "..."}
-
-    Multi-output (DPO / GRPO)::
-
-        {"input": {"sources": {
-            "chosen": {"sections": [{"field": "chosen", "action": "$role", "template": true}]},
-            "rejected": {"sections": [{"field": "rejected", "action": "$role", "template": true}]},
-        }}}
-        → {"chosen": [...], "chosen_mask": [...], "rejected": [...], "rejected_mask": [...], "domain": "..."}
-
-    Output spec fields::
-
-        sections      – list of section specs (same format as single-output)
-        list_field    – True when JSONL field holds a list (GRPO responses)
-        mask_key      – explicit loss-mask output key (default: ``"{output_key}_mask"``)
     """
 
-    def __init__(self):
+    def __init__(self, renderer: Optional[SectionRenderer] = None):
-        self.renderer = SectionRenderer()
+        self.renderer = renderer or SectionRenderer()
 
     def build(self, item: dict, config, tokenizer) -> Optional[dict]:
-        sources_spec = getattr(config.input, "sources", None)
-        if sources_spec:
-            return self._build_multi(item, sources_spec, config, tokenizer)
-        return self._build_single(item, config, tokenizer)
-
-    def _build_single(self, item: dict, config, tokenizer) -> Optional[dict]:
         sections = config.input.sections
         if not sections:
             return None
@@ -266,9 +243,22 @@ class SectionedMaskBuilder(BaseMaskBuilder):
             result["loss_mask"] = mask
         return result
 
-    def _build_multi(
+
-        self, item: dict, sources_spec: dict, config, tokenizer
+@MaskBuilderFactory.register("multi")
-    ) -> Optional[dict]:
+class MultiOutputMaskBuilder(BaseMaskBuilder):
+    """Build multiple output sequences (DPO / GRPO).
+
+    Expects ``config.input.sources`` (dict of output_key → spec).
+    """
+
+    def __init__(self, renderer: Optional[SectionRenderer] = None):
+        self.renderer = renderer or SectionRenderer()
+
+    def build(self, item: dict, config, tokenizer) -> Optional[dict]:
+        sources_spec = getattr(config.input, "sources", None)
+        if not sources_spec:
+            return None
+
         result: dict = {}
         any_output = False
 
@@ -313,3 +303,22 @@ class SectionedMaskBuilder(BaseMaskBuilder):
 
         result["domain"] = _extract_domain(item, config.output.domain_key)
         return result
+
+
+@MaskBuilderFactory.register("sectioned")
+class SectionedMaskBuilder(BaseMaskBuilder):
+    """Façade that dispatches to SingleOutputMaskBuilder or MultiOutputMaskBuilder.
+
+    Preserves backward compatibility for existing configs and code that rely
+    on the ``"sectioned"`` factory name.
+    """
+
+    def __init__(self):
+        self._single = SingleOutputMaskBuilder()
+        self._multi = MultiOutputMaskBuilder()
+
+    def build(self, item: dict, config, tokenizer) -> Optional[dict]:
+        sources_spec = getattr(config.input, "sources", None)
+        if sources_spec:
+            return self._multi.build(item, config, tokenizer)
+        return self._single.build(item, config, tokenizer)

tests/data/conftest.py

@@ -10,7 +10,11 @@ from astrai.config.preprocess_config import (
     PipelineConfig,
     ProcessingConfig,
 )
-from astrai.preprocessing.builder import SectionedMaskBuilder
+from astrai.preprocessing.builder import (
+    MultiOutputMaskBuilder,
+    SectionedMaskBuilder,
+    SingleOutputMaskBuilder,
+)
 from astrai.tokenize import AutoTokenizer
 
 _SPECIAL_TOKENS_CONFIG = {
@@ -210,6 +214,16 @@ def builder():
     return SectionedMaskBuilder()
 
 
+@pytest.fixture
+def single_builder():
+    return SingleOutputMaskBuilder()
+
+
+@pytest.fixture
+def multi_builder():
+    return MultiOutputMaskBuilder()
+
+
 @pytest.fixture
 def tokenizer_dir(temp_dir, test_tokenizer):
     d = os.path.join(temp_dir, "tok")

tests/data/test_preprocess_builder.py

@@ -8,7 +8,9 @@ from astrai.config.preprocess_config import (
 )
 from astrai.preprocessing.builder import (
     MaskBuilderFactory,
+    MultiOutputMaskBuilder,
     SectionedMaskBuilder,
+    SingleOutputMaskBuilder,
 )
 from tests.data.conftest import (
     _CHAT_SECTIONS,
@@ -272,12 +274,18 @@ def test_sectioned_text_too_short(test_tokenizer, builder):
 
 def test_factory_registered():
     names = MaskBuilderFactory.list_registered()
+    assert "single" in names
+    assert "multi" in names
     assert "sectioned" in names
 
 
 def test_factory_create():
-    builder_obj = MaskBuilderFactory.create("sectioned")
+    single = MaskBuilderFactory.create("single")
-    assert isinstance(builder_obj, SectionedMaskBuilder)
+    assert isinstance(single, SingleOutputMaskBuilder)
+    multi = MaskBuilderFactory.create("multi")
+    assert isinstance(multi, MultiOutputMaskBuilder)
+    sectioned = MaskBuilderFactory.create("sectioned")
+    assert isinstance(sectioned, SectionedMaskBuilder)
 
 
 def test_dpo_chat_basic(chat_tokenizer, builder):
@@ -367,3 +375,59 @@ def test_grpo_single_reward(chat_tokenizer, builder):
     }
     result = builder.build(item, config, chat_tokenizer)
     assert result["rewards"] == [0.9]
+
+
+def test_single_builder_matches_facade(chat_tokenizer, builder, single_builder):
+    config = make_chat_config()
+    item = {
+        "messages": [
+            {"role": "user", "content": "What is 2+2?"},
+            {"role": "assistant", "content": "4"},
+        ]
+    }
+    facade_result = builder.build(item, config, chat_tokenizer)
+    single_result = single_builder.build(item, config, chat_tokenizer)
+    assert single_result == facade_result
+
+
+def test_single_builder_rejects_multi_config(chat_tokenizer, single_builder):
+    config = make_dpo_chat_config()
+    item = {
+        "chosen": [
+            {"role": "user", "content": "What is 2+2?"},
+            {"role": "assistant", "content": "4"},
+        ],
+        "rejected": [
+            {"role": "user", "content": "What is 2+2?"},
+            {"role": "assistant", "content": "5"},
+        ],
+    }
+    assert single_builder.build(item, config, chat_tokenizer) is None
+
+
+def test_multi_builder_matches_facade(chat_tokenizer, builder, multi_builder):
+    config = make_dpo_chat_config()
+    item = {
+        "chosen": [
+            {"role": "user", "content": "What is 2+2?"},
+            {"role": "assistant", "content": "4"},
+        ],
+        "rejected": [
+            {"role": "user", "content": "What is 2+2?"},
+            {"role": "assistant", "content": "5"},
+        ],
+    }
+    facade_result = builder.build(item, config, chat_tokenizer)
+    multi_result = multi_builder.build(item, config, chat_tokenizer)
+    assert multi_result == facade_result
+
+
+def test_multi_builder_rejects_single_config(chat_tokenizer, multi_builder):
+    config = make_chat_config()
+    item = {
+        "messages": [
+            {"role": "user", "content": "What is 2+2?"},
+            {"role": "assistant", "content": "4"},
+        ]
+    }
+    assert multi_builder.build(item, config, chat_tokenizer) is None