refactor : 压缩测试代码，消除重复

- fixture 替代重复实例化和 tokenizer 落盘
- parametrize 合并同构测试
- helper 消除 save_h5 + DatasetFactory.load 样板
- 净减 272 行

tests/data/conftest.py

@@ -1,3 +1,5 @@
+import json
+import os
 import tempfile
 
 import pytest
@@ -8,6 +10,7 @@ from astrai.config.preprocess_config import (
     PipelineConfig,
     ProcessingConfig,
 )
+from astrai.preprocessing.builder import SectionedMaskBuilder
 from astrai.tokenize import AutoTokenizer
 
 _SPECIAL_TOKENS_CONFIG = {
@@ -200,3 +203,33 @@ def make_grpo_no_template_config():
         mask_default="mask",
         preprocessing=ProcessingConfig(max_seq_len=2048),
     )
+
+
+@pytest.fixture
+def builder():
+    return SectionedMaskBuilder()
+
+
+@pytest.fixture
+def tokenizer_dir(temp_dir, test_tokenizer):
+    d = os.path.join(temp_dir, "tok")
+    os.makedirs(d, exist_ok=True)
+    test_tokenizer._tokenizer.save(os.path.join(d, "tokenizer.json"))
+    with open(os.path.join(d, "tokenizer_config.json"), "w") as f:
+        json.dump(
+            {"special_tokens": {"pad_token": "<|_pad_|>", "unk_token": "<|_unk_|>"}}, f
+        )
+    return d
+
+
+@pytest.fixture
+def chat_tokenizer_dir(temp_dir, chat_tokenizer):
+    d = os.path.join(temp_dir, "tok")
+    os.makedirs(d, exist_ok=True)
+    chat_tokenizer._tokenizer.save(os.path.join(d, "tokenizer.json"))
+    with open(os.path.join(d, "tokenizer_config.json"), "w") as f:
+        json.dump(
+            {"special_tokens": _SPECIAL_TOKENS_CONFIG, "chat_template": _CHAT_TEMPLATE},
+            f,
+        )
+    return d

tests/data/test_dataset.py

@@ -15,28 +15,34 @@ from astrai.dataset.storage import (
 )
 
 
+def _rand_seq(length, vocab=1000):
+    return torch.randint(0, vocab, (length,), dtype=torch.int64)
+
+
+def _make_seq_dataset(
+    test_dir, name="data", seq_length=200, train_type="seq", data=None, **load_kwargs
+):
+    if data is None:
+        data = {"sequence": [_rand_seq(seq_length)]}
+    save_h5(test_dir, name, data)
+    return DatasetFactory.load(
+        train_type,
+        test_dir,
+        window_size=load_kwargs.pop("window_size", 64),
+        **load_kwargs,
+    )
+
+
 def test_dataset_loader_random_paths(base_test_env):
     """Test dataset loader with multiple random paths"""
     test_dir = base_test_env["test_dir"]
 
-    # Create multiple mmap dataset directories with random data
     num_files = np.random.randint(2, 5)
-
     for i in range(num_files):
         seq_length = np.random.randint(200, 400)
-        dummy_data = {
-            "sequence": [
-                torch.randint(0, 1000, (seq_length,), dtype=torch.int64)
-                for _ in range(10)
-            ],
-        }
-        save_h5(test_dir, f"data_{i}", dummy_data)
-
-        # Test loading with multiple paths
-        loaded_dataset = DatasetFactory.load(
-            train_type="seq",
-            load_path=test_dir,
-            window_size=64,
+        dummy_data = {"sequence": [_rand_seq(seq_length) for _ in range(10)]}
+        loaded_dataset = _make_seq_dataset(
+            test_dir, f"data_{i}", seq_length, data=dummy_data
         )
         assert loaded_dataset is not None
         assert len(loaded_dataset) > 0
@@ -54,23 +60,15 @@ def test_dpo_strategy_with_random_data(base_test_env):
     """Test DPO strategy with randomized preference data"""
     test_dir = base_test_env["test_dir"]
 
-    # Create DPO-style data with memory mapping format
     seq_length = np.random.randint(100, 200)
-
     dummy_data = {
-        "chosen": [torch.randint(0, 1000, (seq_length,), dtype=torch.int64)],
-        "rejected": [torch.randint(0, 1000, (seq_length,), dtype=torch.int64)],
+        "chosen": [_rand_seq(seq_length)],
+        "rejected": [_rand_seq(seq_length)],
         "chosen_mask": [torch.ones(seq_length, dtype=torch.bool)],
         "rejected_mask": [torch.ones(seq_length, dtype=torch.bool)],
     }
-
-    save_h5(test_dir, "dpo_data", dummy_data)
-
-    # Load DPO dataset
-    dpo_dataset = DatasetFactory.load(
-        train_type="dpo",
-        load_path=test_dir,
-        window_size=64,
+    dpo_dataset = _make_seq_dataset(
+        test_dir, "dpo_data", seq_length, train_type="dpo", data=dummy_data
     )
 
     assert dpo_dataset is not None
@@ -92,22 +90,14 @@ def test_sft_dataset_with_random_data(base_test_env):
     """Test SFT dataset with random data"""
     test_dir = base_test_env["test_dir"]
 
-    # Create SFT-style data with memory mapping format
     seq_length = np.random.randint(100, 200)
-
     dummy_data = {
-        "sequence": [torch.randint(0, 1000, (seq_length,), dtype=torch.int64)],
+        "sequence": [_rand_seq(seq_length)],
         "loss_mask": [torch.ones(seq_length, dtype=torch.bool)],
         "position_ids": [torch.arange(seq_length, dtype=torch.int32)],
     }
-
-    save_h5(test_dir, "sft_data", dummy_data)
-
-    # Load SFT dataset
-    sft_dataset = DatasetFactory.load(
-        train_type="sft",
-        load_path=test_dir,
-        window_size=64,
+    sft_dataset = _make_seq_dataset(
+        test_dir, "sft_data", seq_length, train_type="sft", data=dummy_data
     )
 
     assert sft_dataset is not None
@@ -128,25 +118,11 @@ def test_dataset_with_custom_stride(base_test_env):
     """Test dataset with custom stride parameter"""
     test_dir = base_test_env["test_dir"]
 
-    # Create test data
-    seq_length = 200
-    dummy_data = {
-        "sequence": [torch.randint(0, 1000, (seq_length,), dtype=torch.int64)],
-    }
-
-    save_h5(test_dir, "stride_test_data", dummy_data)
-
-    # Test with custom stride
     custom_stride = 32
-    dataset = DatasetFactory.load(
-        train_type="seq", load_path=test_dir, window_size=64, stride=custom_stride
-    )
-
+    dataset = _make_seq_dataset(test_dir, "stride_test_data", stride=custom_stride)
     assert dataset is not None
     assert len(dataset) > 0
 
-    # With stride 32 and window 64 on 200 length data, we should get more samples
-    # than with default stride (which equals window size)
     default_stride_dataset = DatasetFactory.load(
         train_type="seq",
         load_path=test_dir,
@@ -157,25 +133,11 @@ def test_dataset_with_custom_stride(base_test_env):
 
 
 def test_dataset_count_property(base_test_env):
-    """Test the count property returns correct raw token count"""
     test_dir = base_test_env["test_dir"]
-
-    seq_length = 200
-    dummy_data = {
-        "sequence": [torch.randint(0, 1000, (seq_length,), dtype=torch.int64)],
-    }
-
-    save_h5(test_dir, "count_test_data", dummy_data)
-
-    dataset = DatasetFactory.load(
-        train_type="seq",
-        load_path=test_dir,
-        window_size=64,
-    )
-
-    assert dataset.count == seq_length
-    assert dataset.count > len(dataset)  # raw tokens > windows
-    assert len(dataset) == (seq_length - 1 - 64) // 64 + 1
+    dataset = _make_seq_dataset(test_dir, "count_test_data")
+    assert dataset.count == 200
+    assert dataset.count > len(dataset)
+    assert len(dataset) == (200 - 1 - 64) // 64 + 1
 
 
 def test_empty_dataset_count():
@@ -186,17 +148,10 @@ def test_empty_dataset_count():
 
 
 def test_dataset_too_short_for_window(base_test_env):
-    """Dataset shorter than window_size returns __len__ == 0"""
     test_dir = base_test_env["test_dir"]
-    seq_length = 30
-    save_h5(
-        test_dir,
-        "short",
-        {"sequence": [torch.randint(0, 1000, (seq_length,), dtype=torch.int64)]},
-    )
-    dataset = DatasetFactory.load("seq", test_dir, window_size=64)
+    dataset = _make_seq_dataset(test_dir, "short", seq_length=30)
     assert len(dataset) == 0
-    assert dataset.count == seq_length
+    assert dataset.count == 30
 
 
 def test_unloaded_dataset_getitem_raises():
@@ -220,12 +175,8 @@ def test_store_unloaded_len():
 
 
 def test_store_fetch_begin_equals_end(base_test_env):
-    """Store.fetch with begin == end returns empty tensor"""
     test_dir = base_test_env["test_dir"]
-    dummy = {"sequence": [torch.randint(0, 1000, (100,), dtype=torch.int64)]}
-    save_h5(test_dir, "empty_fetch", dummy)
-
-    dataset = DatasetFactory.load("seq", test_dir, window_size=32)
+    dataset = _make_seq_dataset(test_dir, "empty_fetch", seq_length=100, window_size=32)
     result = dataset.storage.fetch(10, 10, "sequence")
     assert result.numel() == 0
 
@@ -299,12 +250,8 @@ def test_save_load_bin_roundtrip(base_test_env):
 
 
 def test_mmap_store_load_and_fetch(base_test_env):
-    """MmapStore loads bin data and fetches correctly"""
     test_dir = base_test_env["test_dir"]
-
-    data = {
-        "sequence": [torch.randint(0, 1000, (200,), dtype=torch.int64)],
-    }
+    data = {"sequence": [_rand_seq(200)]}
     save_bin(test_dir, data)
 
     store = StoreFactory.create("bin")
@@ -317,14 +264,9 @@ def test_mmap_store_load_and_fetch(base_test_env):
 
 
 def test_mmap_dataset_load(base_test_env):
-    """DatasetFactory.load auto-detects bin format"""
     test_dir = base_test_env["test_dir"]
-
-    data = {
-        "sequence": [torch.randint(0, 1000, (200,), dtype=torch.int64)],
-    }
+    data = {"sequence": [_rand_seq(200)]}
     save_bin(test_dir, data)
-
     dataset = DatasetFactory.load("seq", test_dir, window_size=64)
     assert len(dataset) > 0
     assert dataset.count == 200
@@ -348,19 +290,16 @@ def test_normalize_mixed_empty_key():
 
 
 def test_grpo_dataset_dtype(base_test_env):
-    """GRPODataset returns correct dtypes"""
     test_dir = base_test_env["test_dir"]
-
-    seq_len = 100
-    data = {
-        "prompts": [torch.randint(0, 100, (seq_len,), dtype=torch.int32)],
-        "responses": [torch.randint(0, 100, (seq_len,), dtype=torch.int32)],
-        "masks": [torch.ones(seq_len, dtype=torch.int32)],
-        "rewards": [torch.ones(seq_len, dtype=torch.float32)],
+    dummy_data = {
+        "prompts": [torch.randint(0, 100, (100,), dtype=torch.int32)],
+        "responses": [torch.randint(0, 100, (100,), dtype=torch.int32)],
+        "masks": [torch.ones(100, dtype=torch.int32)],
+        "rewards": [torch.ones(100, dtype=torch.float32)],
     }
-    save_h5(test_dir, "grpo_dtype", data)
-
-    dataset = DatasetFactory.load("grpo", test_dir, window_size=32)
+    dataset = _make_seq_dataset(
+        test_dir, "grpo_dtype", train_type="grpo", data=dummy_data, window_size=32
+    )
     item = dataset[0]
 
     assert item["prompts"].dtype == torch.long
@@ -370,18 +309,16 @@ def test_grpo_dataset_dtype(base_test_env):
 
 
 def test_grpo_dataset_load(base_test_env):
-    """GRPODataset loads and returns correct keys"""
     test_dir = base_test_env["test_dir"]
-    seq_len = 200
-    data = {
-        "prompts": [torch.randint(0, 1000, (seq_len,), dtype=torch.int64)],
-        "responses": [torch.randint(0, 1000, (seq_len,), dtype=torch.int64)],
-        "masks": [torch.ones(seq_len, dtype=torch.int64)],
-        "rewards": [torch.rand(seq_len, dtype=torch.float32)],
+    dummy_data = {
+        "prompts": [_rand_seq(200)],
+        "responses": [_rand_seq(200)],
+        "masks": [torch.ones(200, dtype=torch.int64)],
+        "rewards": [torch.rand(200, dtype=torch.float32)],
     }
-    save_h5(test_dir, "grpo_test", data)
-
-    dataset = DatasetFactory.load("grpo", test_dir, window_size=64)
+    dataset = _make_seq_dataset(
+        test_dir, "grpo_test", train_type="grpo", data=dummy_data
+    )
     assert len(dataset) > 0
     item = dataset[0]
     assert "prompts" in item
@@ -400,7 +337,6 @@ def test_detect_format_bin_dir(base_test_env):
 
 
 def test_store_fetch_multi_key(base_test_env):
-    """Store.fetch with List[str] returns Dict[str, Tensor]"""
     test_dir = base_test_env["test_dir"]
     save_h5(
         test_dir,
@@ -410,7 +346,6 @@ def test_store_fetch_multi_key(base_test_env):
             "loss_mask": [torch.ones(100, dtype=torch.int64)],
         },
     )
-
     store = StoreFactory.create("h5")
     store.load(test_dir)
     result = store.fetch(10, 20, ["sequence", "loss_mask"])
@@ -420,10 +355,8 @@ def test_store_fetch_multi_key(base_test_env):
 
 
 def test_store_fetch_out_of_bounds(base_test_env):
-    """Store.fetch raises ValueError for out-of-bounds indices"""
     test_dir = base_test_env["test_dir"]
     save_h5(test_dir, "bounds", {"sequence": [torch.randint(0, 100, (50,))]})
-
     store = StoreFactory.create("h5")
     store.load(test_dir)
     with pytest.raises(ValueError, match="out of bounds"):
@@ -435,10 +368,7 @@ def test_store_fetch_out_of_bounds(base_test_env):
 
 
 def test_dataset_load_explicit_storage_type(base_test_env):
-    """DatasetFactory.load with explicit storage_type bypasses auto-detect"""
     test_dir = base_test_env["test_dir"]
-    save_h5(test_dir, "explicit", {"sequence": [torch.randint(0, 100, (200,))]})
-
-    dataset = DatasetFactory.load("seq", test_dir, window_size=64, storage_type="h5")
+    dataset = _make_seq_dataset(test_dir, "explicit", storage_type="h5")
     assert len(dataset) > 0
     assert dataset.count == 200

tests/data/test_preprocess_builder.py

@@ -1,3 +1,5 @@
+import pytest
+
 from astrai.config.preprocess_config import (
     InputConfig,
     OutputConfig,
@@ -20,9 +22,8 @@ from tests.data.conftest import (
 )
 
 
-def test_chat_simple(chat_tokenizer):
+def test_chat_simple(chat_tokenizer, builder):
     config = make_chat_config()
-    builder = SectionedMaskBuilder()
     item = {
         "messages": [
             {"role": "system", "content": "You are helpful."},
@@ -46,9 +47,8 @@ def test_chat_simple(chat_tokenizer):
     assert trained < total
 
 
-def test_chat_mask_only_assistant(chat_tokenizer):
+def test_chat_mask_only_assistant(chat_tokenizer, builder):
     config = make_chat_config()
-    builder = SectionedMaskBuilder()
     item = {
         "messages": [
             {"role": "user", "content": "What is 2+2?"},
@@ -66,14 +66,22 @@ def test_chat_mask_only_assistant(chat_tokenizer):
     assert len(masked) > 0
 
 
-def test_chat_all_masked(chat_tokenizer):
+@pytest.mark.parametrize(
+    "mask_rules,mask_default,expect_nonzero",
+    [
+        ({"system": "mask", "user": "mask", "assistant": "mask"}, "mask", False),
+        ({}, "train", True),
+    ],
+)
+def test_chat_uniform_masking(
+    mask_rules, mask_default, expect_nonzero, chat_tokenizer, builder
+):
     config = PipelineConfig(
         input=InputConfig(sections=_CHAT_SECTIONS),
-        mask={"system": "mask", "user": "mask", "assistant": "mask"},
-        mask_default="mask",
+        mask=mask_rules,
+        mask_default=mask_default,
         preprocessing=ProcessingConfig(max_seq_len=2048),
     )
-    builder = SectionedMaskBuilder()
     item = {
         "messages": [
             {"role": "system", "content": "You are helpful."},
@@ -81,35 +89,20 @@ def test_chat_all_masked(chat_tokenizer):
         ]
     }
     result = builder.build(item, config, chat_tokenizer)
-    assert sum(result["loss_mask"]) == 0
+    masked_count = sum(result["loss_mask"])
+    if expect_nonzero:
+        assert masked_count > 0
+    else:
+        assert masked_count == 0
 
 
-def test_chat_all_trained(chat_tokenizer):
-    config = PipelineConfig(
-        input=InputConfig(sections=_CHAT_SECTIONS),
-        mask={},
-        mask_default="train",
-        preprocessing=ProcessingConfig(max_seq_len=2048),
-    )
-    builder = SectionedMaskBuilder()
-    item = {
-        "messages": [
-            {"role": "system", "content": "You are helpful."},
-            {"role": "assistant", "content": "Hi there!"},
-        ]
-    }
-    result = builder.build(item, config, chat_tokenizer)
-    assert sum(result["loss_mask"]) == len(result["sequence"]) - 1
-
-
-def test_chat_empty_messages(chat_tokenizer):
+def test_chat_empty_messages(chat_tokenizer, builder):
     config = make_chat_config()
-    builder = SectionedMaskBuilder()
     assert builder.build({"messages": []}, config, chat_tokenizer) is None
     assert builder.build({}, config, chat_tokenizer) is None
 
 
-def test_chat_domain_extraction(chat_tokenizer):
+def test_chat_domain_extraction(chat_tokenizer, builder):
     config = PipelineConfig(
         input=InputConfig(sections=_CHAT_SECTIONS),
         mask={"assistant": "train"},
@@ -117,7 +110,6 @@ def test_chat_domain_extraction(chat_tokenizer):
         preprocessing=ProcessingConfig(max_seq_len=2048),
         output=OutputConfig(domain_key="source"),
     )
-    builder = SectionedMaskBuilder()
     item = {
         "messages": [
             {"role": "user", "content": "Hi"},
@@ -129,14 +121,13 @@ def test_chat_domain_extraction(chat_tokenizer):
     assert result["domain"] == "wiki"
 
 
-def test_chat_truncation(chat_tokenizer):
+def test_chat_truncation(chat_tokenizer, builder):
     config = PipelineConfig(
         input=InputConfig(sections=_CHAT_SECTIONS),
         mask={"assistant": "train"},
         mask_default="mask",
         preprocessing=ProcessingConfig(max_seq_len=10),
     )
-    builder = SectionedMaskBuilder()
     item = {
         "messages": [
             {
@@ -151,18 +142,16 @@ def test_chat_truncation(chat_tokenizer):
     assert len(result["loss_mask"]) == len(result["sequence"])
 
 
-def test_instruction_basic(test_tokenizer):
+def test_instruction_basic(test_tokenizer, builder):
     config = make_instruction_config()
-    builder = SectionedMaskBuilder()
     item = {"prompt": "Translate to French: Hello", "response": "Bonjour"}
     result = builder.build(item, config, test_tokenizer)
     assert result is not None
     assert len(result["sequence"]) == len(result["loss_mask"])
 
 
-def test_instruction_prompt_masked(test_tokenizer):
+def test_instruction_prompt_masked(test_tokenizer, builder):
     config = make_instruction_config()
-    builder = SectionedMaskBuilder()
     item = {"prompt": "hello", "response": "world"}
     result = builder.build(item, config, test_tokenizer)
     mask = result["loss_mask"]
@@ -175,7 +164,7 @@ def test_instruction_prompt_masked(test_tokenizer):
         assert all(m == 1 for m in mask[p_len:])
 
 
-def test_instruction_train_on_prompt(test_tokenizer):
+def test_instruction_train_on_prompt(test_tokenizer, builder):
     config = PipelineConfig(
         input=InputConfig(
             sections=[
@@ -185,7 +174,6 @@ def test_instruction_train_on_prompt(test_tokenizer):
         ),
         preprocessing=ProcessingConfig(max_seq_len=2048),
     )
-    builder = SectionedMaskBuilder()
     item = {"prompt": "hello", "response": "world"}
     result = builder.build(item, config, test_tokenizer)
     mask = result["loss_mask"]
@@ -196,9 +184,8 @@ def test_instruction_train_on_prompt(test_tokenizer):
     assert all(m == 1 for m in mask[:p_len])
 
 
-def test_text_basic(test_tokenizer):
+def test_text_basic(test_tokenizer, builder):
     config = make_text_config()
-    builder = SectionedMaskBuilder()
     item = {"text": "Hello world. This is a test document."}
     result = builder.build(item, config, test_tokenizer)
     assert result is not None
@@ -207,41 +194,37 @@ def test_text_basic(test_tokenizer):
     assert "loss_mask" not in result
 
 
-def test_text_empty(test_tokenizer):
+def test_text_empty(test_tokenizer, builder):
     config = make_text_config()
-    builder = SectionedMaskBuilder()
     assert builder.build({"text": ""}, config, test_tokenizer) is None
     assert builder.build({"text": "   "}, config, test_tokenizer) is None
 
 
-def test_text_too_short(test_tokenizer):
+def test_text_too_short(test_tokenizer, builder):
     config = PipelineConfig(
         input=InputConfig(sections=_TEXT_SECTIONS),
         preprocessing=ProcessingConfig(min_chars=100),
     )
-    builder = SectionedMaskBuilder()
     assert builder.build({"text": "short"}, config, test_tokenizer) is None
 
 
-def test_text_truncation(test_tokenizer):
+def test_text_truncation(test_tokenizer, builder):
     config = PipelineConfig(
         input=InputConfig(sections=_TEXT_SECTIONS),
         preprocessing=ProcessingConfig(max_seq_len=3, min_chars=1),
     )
-    builder = SectionedMaskBuilder()
     item = {"text": "This is a very long text that should be truncated"}
     result = builder.build(item, config, test_tokenizer)
     assert len(result["sequence"]) <= 3
 
 
-def test_sectioned_chat(chat_tokenizer):
+def test_sectioned_chat(chat_tokenizer, builder):
     config = PipelineConfig(
         input=InputConfig(sections=_CHAT_SECTIONS),
         mask={"system": "mask", "user": "mask", "assistant": "train"},
         mask_default="mask",
         preprocessing=ProcessingConfig(max_seq_len=2048),
     )
-    builder = SectionedMaskBuilder()
     item = {
         "messages": [
             {"role": "user", "content": "What is 2+2?"},
@@ -255,12 +238,11 @@ def test_sectioned_chat(chat_tokenizer):
     assert 0 in result["loss_mask"]
 
 
-def test_sectioned_instruction(test_tokenizer):
+def test_sectioned_instruction(test_tokenizer, builder):
     config = PipelineConfig(
         input=InputConfig(sections=_INSTRUCTION_SECTIONS),
         preprocessing=ProcessingConfig(max_seq_len=2048, min_chars=0),
     )
-    builder = SectionedMaskBuilder()
     item = {"prompt": "Q: Why?", "response": "A: Because."}
     result = builder.build(item, config, test_tokenizer)
     assert result is not None
@@ -269,24 +251,22 @@ def test_sectioned_instruction(test_tokenizer):
     assert mask[-1] == 1
 
 
-def test_sectioned_text(test_tokenizer):
+def test_sectioned_text(test_tokenizer, builder):
     config = PipelineConfig(
         input=InputConfig(sections=_TEXT_SECTIONS),
         preprocessing=ProcessingConfig(max_seq_len=2048, min_chars=1),
     )
-    builder = SectionedMaskBuilder()
     item = {"text": "Hello world, this is a test."}
     result = builder.build(item, config, test_tokenizer)
     assert result is not None
     assert "loss_mask" not in result
 
 
-def test_sectioned_text_too_short(test_tokenizer):
+def test_sectioned_text_too_short(test_tokenizer, builder):
     config = PipelineConfig(
         input=InputConfig(sections=_TEXT_SECTIONS),
         preprocessing=ProcessingConfig(max_seq_len=2048, min_chars=100),
     )
-    builder = SectionedMaskBuilder()
     assert builder.build({"text": "short"}, config, test_tokenizer) is None
 
 
@@ -296,13 +276,12 @@ def test_factory_registered():
 
 
 def test_factory_create():
-    builder = MaskBuilderFactory.create("sectioned")
-    assert isinstance(builder, SectionedMaskBuilder)
+    builder_obj = MaskBuilderFactory.create("sectioned")
+    assert isinstance(builder_obj, SectionedMaskBuilder)
 
 
-def test_dpo_chat_basic(chat_tokenizer):
+def test_dpo_chat_basic(chat_tokenizer, builder):
     config = make_dpo_chat_config()
-    builder = SectionedMaskBuilder()
     item = {
         "chosen": [
             {"role": "user", "content": "What is 2+2?"},
@@ -319,16 +298,14 @@ def test_dpo_chat_basic(chat_tokenizer):
     assert "rejected" in result
     assert "chosen_mask" in result
     assert "rejected_mask" in result
-    assert "domain" in result
     assert len(result["chosen"]) == len(result["chosen_mask"])
     assert len(result["rejected"]) == len(result["rejected_mask"])
     assert sum(result["chosen_mask"]) > 0
     assert sum(result["rejected_mask"]) > 0
 
 
-def test_dpo_chosen_only_trained(chat_tokenizer):
+def test_dpo_chosen_only_trained(chat_tokenizer, builder):
     config = make_dpo_chat_config()
-    builder = SectionedMaskBuilder()
     item = {
         "chosen": [
             {"role": "user", "content": "Hi"},
@@ -346,15 +323,13 @@ def test_dpo_chosen_only_trained(chat_tokenizer):
     assert 1 in result["rejected_mask"]
 
 
-def test_dpo_missing_field_is_none(chat_tokenizer):
+def test_dpo_missing_field_is_none(chat_tokenizer, builder):
     config = make_dpo_chat_config()
-    builder = SectionedMaskBuilder()
     assert builder.build({"chosen": [], "rejected": []}, config, chat_tokenizer) is None
 
 
-def test_grpo_basic(chat_tokenizer):
+def test_grpo_basic(chat_tokenizer, builder):
     config = make_grpo_config()
-    builder = SectionedMaskBuilder()
     item = {
         "prompt": [{"role": "user", "content": "What is 2+2?"}],
         "responses": ["4", "The answer is four", "Four", "2+2=4"],
@@ -370,9 +345,8 @@ def test_grpo_basic(chat_tokenizer):
     assert result["rewards"] == [1.0, 0.5, 0.8, 0.2]
 
 
-def test_grpo_response_tokens_all_trained(chat_tokenizer):
+def test_grpo_response_tokens_all_trained(chat_tokenizer, builder):
     config = make_grpo_config()
-    builder = SectionedMaskBuilder()
     item = {
         "prompt": [{"role": "user", "content": "Q"}],
         "responses": ["A", "B"],
@@ -384,9 +358,8 @@ def test_grpo_response_tokens_all_trained(chat_tokenizer):
     assert len(masks) == len(result["responses"])
 
 
-def test_grpo_single_reward(chat_tokenizer):
+def test_grpo_single_reward(chat_tokenizer, builder):
     config = make_grpo_config()
-    builder = SectionedMaskBuilder()
     item = {
         "prompt": [{"role": "user", "content": "Q"}],
         "responses": ["A"],

tests/data/test_preprocess_pipeline.py

@@ -10,9 +10,7 @@ from astrai.config.preprocess_config import (
 from astrai.preprocessing.pipeline import Pipeline, filter_by_length
 from tests.data.conftest import (
     _CHAT_SECTIONS,
-    _CHAT_TEMPLATE,
     _INSTRUCTION_SECTIONS,
-    _SPECIAL_TOKENS_CONFIG,
     _TEXT_SECTIONS,
     make_dpo_chat_config,
     make_grpo_no_template_config,
@@ -26,19 +24,7 @@ def test_filter_by_length():
     assert filter_by_length("just right", min_len=5, max_len=20)
 
 
-def test_full_chat_pipeline(temp_dir, chat_tokenizer):
-    tokenizer_dir = os.path.join(temp_dir, "tok")
-    os.makedirs(tokenizer_dir, exist_ok=True)
-    chat_tokenizer._tokenizer.save(os.path.join(tokenizer_dir, "tokenizer.json"))
-    with open(os.path.join(tokenizer_dir, "tokenizer_config.json"), "w") as f:
-        json.dump(
-            {
-                "special_tokens": _SPECIAL_TOKENS_CONFIG,
-                "chat_template": _CHAT_TEMPLATE,
-            },
-            f,
-        )
-
+def test_full_chat_pipeline(temp_dir, chat_tokenizer_dir):
     jsonl_path = os.path.join(temp_dir, "chat.jsonl")
     with open(jsonl_path, "w", encoding="utf-8") as f:
         f.write(
@@ -78,7 +64,7 @@ def test_full_chat_pipeline(temp_dir, chat_tokenizer):
         config=config,
         input_paths=[jsonl_path],
         output_dir=out_dir,
-        tokenizer_path=tokenizer_dir,
+        tokenizer_path=chat_tokenizer_dir,
     ).run()
 
     meta_path = os.path.join(out_dir, "__default__", "shard_0000", "meta.json")
@@ -91,21 +77,7 @@ def test_full_chat_pipeline(temp_dir, chat_tokenizer):
     assert meta["loss_mask"]["dtype"] == "int32"
 
 
-def test_full_text_pipeline(temp_dir, test_tokenizer):
-    tokenizer_dir = os.path.join(temp_dir, "tok")
-    os.makedirs(tokenizer_dir, exist_ok=True)
-    test_tokenizer._tokenizer.save(os.path.join(tokenizer_dir, "tokenizer.json"))
-    with open(os.path.join(tokenizer_dir, "tokenizer_config.json"), "w") as f:
-        json.dump(
-            {
-                "special_tokens": {
-                    "pad_token": "<|_pad_|>",
-                    "unk_token": "<|_unk_|>",
-                }
-            },
-            f,
-        )
-
+def test_full_text_pipeline(temp_dir, tokenizer_dir):
     jsonl_path = os.path.join(temp_dir, "text.jsonl")
     with open(jsonl_path, "w", encoding="utf-8") as f:
         f.write(
@@ -145,24 +117,9 @@ def test_full_text_pipeline(temp_dir, test_tokenizer):
         meta = json.load(f)
     assert "sequence" in meta
     assert "loss_mask" not in meta
-    assert meta["sequence"]["dtype"] == "int32"
 
 
-def test_full_instruction_pipeline(temp_dir, test_tokenizer):
-    tokenizer_dir = os.path.join(temp_dir, "tok")
-    os.makedirs(tokenizer_dir, exist_ok=True)
-    test_tokenizer._tokenizer.save(os.path.join(tokenizer_dir, "tokenizer.json"))
-    with open(os.path.join(tokenizer_dir, "tokenizer_config.json"), "w") as f:
-        json.dump(
-            {
-                "special_tokens": {
-                    "pad_token": "<|_pad_|>",
-                    "unk_token": "<|_unk_|>",
-                }
-            },
-            f,
-        )
-
+def test_full_instruction_pipeline(temp_dir, tokenizer_dir):
     jsonl_path = os.path.join(temp_dir, "instruct.jsonl")
     with open(jsonl_path, "w", encoding="utf-8") as f:
         f.write(
@@ -206,25 +163,9 @@ def test_full_instruction_pipeline(temp_dir, test_tokenizer):
         meta = json.load(f)
     assert "sequence" in meta
     assert "loss_mask" in meta
-    assert meta["sequence"]["dtype"] == "int32"
-    assert meta["loss_mask"]["dtype"] == "int32"
 
 
-def test_dtype_override(temp_dir, test_tokenizer):
-    tokenizer_dir = os.path.join(temp_dir, "tok")
-    os.makedirs(tokenizer_dir, exist_ok=True)
-    test_tokenizer._tokenizer.save(os.path.join(tokenizer_dir, "tokenizer.json"))
-    with open(os.path.join(tokenizer_dir, "tokenizer_config.json"), "w") as f:
-        json.dump(
-            {
-                "special_tokens": {
-                    "pad_token": "<|_pad_|>",
-                    "unk_token": "<|_unk_|>",
-                }
-            },
-            f,
-        )
-
+def test_dtype_override(temp_dir, tokenizer_dir):
     jsonl_path = os.path.join(temp_dir, "data.jsonl")
     with open(jsonl_path, "w", encoding="utf-8") as f:
         f.write(json.dumps({"prompt": "Q", "response": "A"}) + "\n")
@@ -252,19 +193,7 @@ def test_dtype_override(temp_dir, test_tokenizer):
     assert meta["loss_mask"]["dtype"] == "bool"
 
 
-def test_dpo_pipeline(temp_dir, chat_tokenizer):
-    tokenizer_dir = os.path.join(temp_dir, "tok")
-    os.makedirs(tokenizer_dir, exist_ok=True)
-    chat_tokenizer._tokenizer.save(os.path.join(tokenizer_dir, "tokenizer.json"))
-    with open(os.path.join(tokenizer_dir, "tokenizer_config.json"), "w") as f:
-        json.dump(
-            {
-                "special_tokens": _SPECIAL_TOKENS_CONFIG,
-                "chat_template": _CHAT_TEMPLATE,
-            },
-            f,
-        )
-
+def test_dpo_pipeline(temp_dir, chat_tokenizer_dir):
     jsonl_path = os.path.join(temp_dir, "dpo.jsonl")
     with open(jsonl_path, "w", encoding="utf-8") as f:
         f.write(
@@ -288,7 +217,7 @@ def test_dpo_pipeline(temp_dir, chat_tokenizer):
         config=make_dpo_chat_config(),
         input_paths=[jsonl_path],
         output_dir=out_dir,
-        tokenizer_path=tokenizer_dir,
+        tokenizer_path=chat_tokenizer_dir,
     ).run()
 
     meta_path = os.path.join(out_dir, "__default__", "shard_0000", "meta.json")
@@ -302,21 +231,7 @@ def test_dpo_pipeline(temp_dir, chat_tokenizer):
     assert "sequence" not in meta
 
 
-def test_grpo_pipeline(temp_dir, test_tokenizer):
-    tokenizer_dir = os.path.join(temp_dir, "tok")
-    os.makedirs(tokenizer_dir, exist_ok=True)
-    test_tokenizer._tokenizer.save(os.path.join(tokenizer_dir, "tokenizer.json"))
-    with open(os.path.join(tokenizer_dir, "tokenizer_config.json"), "w") as f:
-        json.dump(
-            {
-                "special_tokens": {
-                    "pad_token": "<|_pad_|>",
-                    "unk_token": "<|_unk_|>",
-                }
-            },
-            f,
-        )
-
+def test_grpo_pipeline(temp_dir, tokenizer_dir):
     jsonl_path = os.path.join(temp_dir, "grpo.jsonl")
     with open(jsonl_path, "w", encoding="utf-8") as f:
         f.write(