diff --git a/astrai/preprocessing/__init__.py b/astrai/preprocessing/__init__.py index 7d4d6b5..349e7e5 100644 --- a/astrai/preprocessing/__init__.py +++ b/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", diff --git a/astrai/preprocessing/builder.py b/astrai/preprocessing/builder.py index 936dd8f..b798e3e 100644 --- a/astrai/preprocessing/builder.py +++ b/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` -orchestrates single-output / multi-output (DPO / GRPO) assembly. +masks (template / text / value extraction). :class:`SingleOutputMaskBuilder` +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") -class SectionedMaskBuilder(BaseMaskBuilder): - """Config-driven builder supporting single and multi-output modes. +@MaskBuilderFactory.register("single") +class SingleOutputMaskBuilder(BaseMaskBuilder): + """Build a single output sequence with optional loss mask. - Single-output:: - - {"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"``) + Expects ``config.input.sections`` (list of section specs). """ - def __init__(self): - self.renderer = SectionRenderer() + 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 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 - ) -> Optional[dict]: + +@MaskBuilderFactory.register("multi") +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) diff --git a/tests/data/conftest.py b/tests/data/conftest.py index cfb25c7..848ac9c 100644 --- a/tests/data/conftest.py +++ b/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") diff --git a/tests/data/test_preprocess_builder.py b/tests/data/test_preprocess_builder.py index 7f886f0..03a6fc7 100644 --- a/tests/data/test_preprocess_builder.py +++ b/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") - assert isinstance(builder_obj, SectionedMaskBuilder) + single = MaskBuilderFactory.create("single") + 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