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
This commit is contained in:
+15
-1
@@ -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")
|
||||
|
||||
Reference in New Issue
Block a user