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 = {
+        dummy_data = {"sequence": [_rand_seq(seq_length) for _ in range(10)]}
-            "sequence": [
+        loaded_dataset = _make_seq_dataset(
-                torch.randint(0, 1000, (seq_length,), dtype=torch.int64)
+            test_dir, f"data_{i}", seq_length, data=dummy_data
-                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,
         )
         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)],
+        "chosen": [_rand_seq(seq_length)],
-        "rejected": [torch.randint(0, 1000, (seq_length,), dtype=torch.int64)],
+        "rejected": [_rand_seq(seq_length)],
         "chosen_mask": [torch.ones(seq_length, dtype=torch.bool)],
         "rejected_mask": [torch.ones(seq_length, dtype=torch.bool)],
     }
-
+    dpo_dataset = _make_seq_dataset(
-    save_h5(test_dir, "dpo_data", dummy_data)
+        test_dir, "dpo_data", seq_length, train_type="dpo", data=dummy_data
-
-    # Load DPO dataset
-    dpo_dataset = DatasetFactory.load(
-        train_type="dpo",
-        load_path=test_dir,
-        window_size=64,
     )
 
     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)],
     }
-
+    sft_dataset = _make_seq_dataset(
-    save_h5(test_dir, "sft_data", dummy_data)
+        test_dir, "sft_data", seq_length, train_type="sft", data=dummy_data
-
-    # Load SFT dataset
-    sft_dataset = DatasetFactory.load(
-        train_type="sft",
-        load_path=test_dir,
-        window_size=64,
     )
 
     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(
+    dataset = _make_seq_dataset(test_dir, "stride_test_data", stride=custom_stride)
-        train_type="seq", load_path=test_dir, window_size=64, 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"]
-
+    dataset = _make_seq_dataset(test_dir, "count_test_data")
-    seq_length = 200
+    assert dataset.count == 200
-    dummy_data = {
+    assert dataset.count > len(dataset)
-        "sequence": [torch.randint(0, 1000, (seq_length,), dtype=torch.int64)],
+    assert len(dataset) == (200 - 1 - 64) // 64 + 1
-    }
-
-    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
 
 
 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
+    dataset = _make_seq_dataset(test_dir, "short", 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)
     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)]}
+    dataset = _make_seq_dataset(test_dir, "empty_fetch", seq_length=100, window_size=32)
-    save_h5(test_dir, "empty_fetch", dummy)
-
-    dataset = DatasetFactory.load("seq", test_dir, 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": [_rand_seq(200)]}
-    data = {
-        "sequence": [torch.randint(0, 1000, (200,), dtype=torch.int64)],
-    }
     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": [_rand_seq(200)]}
-    data = {
-        "sequence": [torch.randint(0, 1000, (200,), dtype=torch.int64)],
-    }
     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"]
-
+    dummy_data = {
-    seq_len = 100
+        "prompts": [torch.randint(0, 100, (100,), dtype=torch.int32)],
-    data = {
+        "responses": [torch.randint(0, 100, (100,), dtype=torch.int32)],
-        "prompts": [torch.randint(0, 100, (seq_len,), dtype=torch.int32)],
+        "masks": [torch.ones(100, dtype=torch.int32)],
-        "responses": [torch.randint(0, 100, (seq_len,), dtype=torch.int32)],
+        "rewards": [torch.ones(100, dtype=torch.float32)],
-        "masks": [torch.ones(seq_len, dtype=torch.int32)],
-        "rewards": [torch.ones(seq_len, dtype=torch.float32)],
     }
-    save_h5(test_dir, "grpo_dtype", data)
+    dataset = _make_seq_dataset(
-
+        test_dir, "grpo_dtype", train_type="grpo", data=dummy_data, window_size=32
-    dataset = DatasetFactory.load("grpo", test_dir, 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
+    dummy_data = {
-    data = {
+        "prompts": [_rand_seq(200)],
-        "prompts": [torch.randint(0, 1000, (seq_len,), dtype=torch.int64)],
+        "responses": [_rand_seq(200)],
-        "responses": [torch.randint(0, 1000, (seq_len,), dtype=torch.int64)],
+        "masks": [torch.ones(200, dtype=torch.int64)],
-        "masks": [torch.ones(seq_len, dtype=torch.int64)],
+        "rewards": [torch.rand(200, dtype=torch.float32)],
-        "rewards": [torch.rand(seq_len, dtype=torch.float32)],
     }
-    save_h5(test_dir, "grpo_test", data)
+    dataset = _make_seq_dataset(
-
+        test_dir, "grpo_test", train_type="grpo", data=dummy_data
-    dataset = DatasetFactory.load("grpo", test_dir, window_size=64)
+    )
     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 = _make_seq_dataset(test_dir, "explicit", storage_type="h5")
-
-    dataset = DatasetFactory.load("seq", test_dir, window_size=64, 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=mask_rules,
-        mask_default="mask",
+        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):
+def test_chat_empty_messages(chat_tokenizer, builder):
-    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):
     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")
+    builder_obj = MaskBuilderFactory.create("sectioned")
-    assert isinstance(builder, SectionedMaskBuilder)
+    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):
+def test_full_chat_pipeline(temp_dir, chat_tokenizer_dir):
-    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,
-        )
-
     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):
+def test_full_text_pipeline(temp_dir, tokenizer_dir):
-    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,
-        )
-
     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):
+def test_full_instruction_pipeline(temp_dir, tokenizer_dir):
-    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,
-        )
-
     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):
+def test_dtype_override(temp_dir, tokenizer_dir):
-    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,
-        )
-
     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):
+def test_dpo_pipeline(temp_dir, chat_tokenizer_dir):
-    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,
-        )
-
     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):
+def test_grpo_pipeline(temp_dir, tokenizer_dir):
-    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,
-        )
-
     jsonl_path = os.path.join(temp_dir, "grpo.jsonl")
     with open(jsonl_path, "w", encoding="utf-8") as f:
         f.write(

tests/inference/test_tool_parser.py

@@ -13,67 +13,28 @@ from astrai.inference.api.tool_parser import (
 )
 
 
-def test_scan_complete_simple():
+@pytest.mark.parametrize(
-    end, complete = _scan_json('{"key": "value"}', 0)
+    "text,expected_complete,check_end_eq_len",
-    assert complete is True
+    [
-    assert end == len('{"key": "value"}')
+        ('{"key": "value"}', True, True),
-
+        ('{"outer": {"inner": 1}}', True, True),
-
+        ('{"key": "value"', False, False),
-def test_scan_complete_nested():
+        ('{"outer": {"inner": 1}', False, False),
-    text = '{"outer": {"inner": 1}}'
+        ('{"key": "a{b}c"} extra', True, False),
+        (r'{"key": "a\"b"}', True, False),
+        ('{"a": {"b": {"c": {"d": {"e": 5}}}}}', True, True),
+        ('{"items": [{"x": 1}, {"x": 2}]}', True, True),
+        ('{"fn": "function() { return 1; }"}', True, False),
+        ('{"key": "\u5317\u4eac"}', True, False),
+    ],
+)
+def test_scan_json(text, expected_complete, check_end_eq_len):
     end, complete = _scan_json(text, 0)
-    assert complete is True
+    assert complete is expected_complete
+    if check_end_eq_len:
         assert end == len(text)
 
 
-def test_scan_incomplete_unclosed():
-    end, complete = _scan_json('{"key": "value"', 0)
-    assert complete is False
-
-
-def test_scan_incomplete_nested():
-    end, complete = _scan_json('{"outer": {"inner": 1}', 0)
-    assert complete is False
-
-
-def test_scan_string_braces_ignored():
-    text = '{"key": "a{b}c"} extra'
-    end, complete = _scan_json(text, 0)
-    assert complete is True
-
-
-def test_scan_escaped_quote_ignored():
-    text = r'{"key": "a\"b"}'
-    end, complete = _scan_json(text, 0)
-    assert complete is True
-
-
-def test_scan_deeply_nested():
-    text = '{"a": {"b": {"c": {"d": {"e": 5}}}}}'
-    end, complete = _scan_json(text, 0)
-    assert complete is True
-    assert end == len(text)
-
-
-def test_scan_array_with_braces():
-    text = '{"items": [{"x": 1}, {"x": 2}]}'
-    end, complete = _scan_json(text, 0)
-    assert complete is True
-    assert end == len(text)
-
-
-def test_scan_code_in_string():
-    text = '{"fn": "function() { return 1; }"}'
-    end, complete = _scan_json(text, 0)
-    assert complete is True
-
-
-def test_scan_unicode_chars():
-    text = '{"key": "\u5317\u4eac"}'
-    end, complete = _scan_json(text, 0)
-    assert complete is True
-
-
 def test_find_single_tool_call():
     text = '{"name": "get_weather", "arguments": {"city": "Beijing"}}'
     results = _find_tool_calls(text)
@@ -141,10 +102,7 @@ def test_find_arguments_with_array():
 
 
 def test_find_arguments_with_nested_array_of_objects():
-    text = (
+    text = '{"name": "batch", "arguments": {"rows": [{"id": 1, "val": "a"}, {"id": 2, "val": "b"}]}}'
-        '{"name": "batch", '
-        '"arguments": {"rows": [{"id": 1, "val": "a"}, {"id": 2, "val": "b"}]}}'
-    )
     results = _find_tool_calls(text)
     assert len(results) == 1
     assert '"rows"' in results[0]["args"]
@@ -206,38 +164,26 @@ def test_find_extracts_correct_arg_start_position():
     assert json_str == text
 
 
-def test_partial_with_name():
+@pytest.mark.parametrize(
-    result = _find_partial_tool_call('{"name": "func", "arguments": {"city"')
+    "text,expected_name,expected_complete",
+    [
+        ('{"name": "func", "arguments": {"city"', "func", False),
+        ('{"name": "func", "arguments": {"city": "BJ"}}', "func", None),
+        ("plain text", None, None),
+        ('{"nam', None, None),
+        ('{"name": "deep", "arguments": {"a": {"b": {"c": ', "deep", None),
+        ('{"name": "batch", "arguments": {"items": [1, 2, ', "batch", None),
+    ],
+)
+def test_find_partial_tool_call(text, expected_name, expected_complete):
+    result = _find_partial_tool_call(text)
+    if expected_name is None:
+        assert result is None
+    else:
         assert result is not None
-    assert result["name"] == "func"
+        assert result["name"] == expected_name
-    assert result["complete"] is False
+        if expected_complete is not None:
-
+            assert result["complete"] is expected_complete
-
-def test_partial_with_full_args():
-    result = _find_partial_tool_call('{"name": "func", "arguments": {"city": "BJ"}}')
-    assert result is not None
-    assert result["name"] == "func"
-
-
-def test_partial_no_match():
-    assert _find_partial_tool_call("plain text") is None
-
-
-def test_partial_no_name_yet():
-    assert _find_partial_tool_call('{"nam') is None
-
-
-def test_partial_deeply_nested():
-    result = _find_partial_tool_call('{"name": "deep", "arguments": {"a": {"b": {"c": ')
-    assert result is not None
-    assert result["name"] == "deep"
-    assert '"a"' in result["args"]
-
-
-def test_partial_array_incomplete():
-    result = _find_partial_tool_call('{"name": "batch", "arguments": {"items": [1, 2, ')
-    assert result is not None
-    assert result["name"] == "batch"
 
 
 def test_feed_plain_text():
@@ -269,7 +215,6 @@ def test_feed_tool_call_args_streaming():
     parser = SimpleJsonToolParser()
     d1 = parser.feed('{"name": "f", "arguments": {"x":')
     d2 = parser.feed('{"name": "f", "arguments": {"x": "1"}}')
-
     args_deltas = [
         d
         for batch in (d1, d2)
@@ -332,17 +277,6 @@ def test_feed_content_after_tool_call_is_not_emitted():
     assert parser.has_tool_calls
 
 
-def _collect_args_deltas(parser):
-    args_parts = []
-    for d in parser.feed(parser._text_buffer):
-        if "tool_calls" in d:
-            for tc in d["tool_calls"]:
-                fn = tc.get("function", {})
-                if "arguments" in fn and fn["arguments"]:
-                    args_parts.append(fn["arguments"])
-    return args_parts
-
-
 def _simulate_streaming(parser, text):
     all_delta_names = []
     all_args_chunks = []
@@ -447,7 +381,6 @@ def test_streaming_args_diff_only_emits_new_bytes():
     parser = SimpleJsonToolParser()
     step1 = parser.feed('{"name": "f", "arguments": {"city": "Bei')
     step2 = parser.feed('{"name": "f", "arguments": {"city": "Beijing"}}')
-
     all_args = []
     for step in (step1, step2):
         for d in step:
@@ -500,31 +433,21 @@ def test_parse_complete_with_content():
 
 def test_parse_complete_multiple_tool_calls():
     parser = SimpleJsonToolParser()
-    body = (
+    body = '{"name": "get_weather", "arguments": {"city": "Beijing"}}{"name": "get_time", "arguments": {"tz": "Asia/Shanghai"}}'
-        '{"name": "get_weather", "arguments": {"city": "Beijing"}}'
-        '{"name": "get_time", "arguments": {"tz": "Asia/Shanghai"}}'
-    )
     result = parser.parse_complete(body)
     assert result is not None
     assert len(result["tool_calls"]) == 2
     assert result["tool_calls"][0]["function"]["name"] == "get_weather"
     assert result["tool_calls"][1]["function"]["name"] == "get_time"
-    assert "Beijing" in result["tool_calls"][0]["function"]["arguments"]
-    assert "Asia/Shanghai" in result["tool_calls"][1]["function"]["arguments"]
 
 
 def test_parse_complete_complex_real_world():
     parser = SimpleJsonToolParser()
     body = (
-        '{"name": "send_email", '
+        '{"name": "send_email", "arguments": {'
-        '"arguments": {'
+        '"to": ["a@b.com", "c@d.com"], "cc": null, '
-        '"to": ["a@b.com", "c@d.com"], '
+        '"subject": "Hello World", "body": "This is a test email.", '
-        '"cc": null, '
+        '"priority": 1, "attachments": false}}'
-        '"subject": "Hello World", '
-        '"body": "This is a test email.", '
-        '"priority": 1, '
-        '"attachments": false'
-        "}}"
     )
     result = parser.parse_complete(body)
     assert result is not None
@@ -539,11 +462,7 @@ def test_parse_complete_complex_real_world():
 
 def test_parse_complete_content_with_multiple_tool_calls():
     parser = SimpleJsonToolParser()
-    body = (
+    body = 'I will do two things. {"name": "f1", "arguments": {"a": 1}}{"name": "f2", "arguments": {"b": 2}}'
-        "I will do two things. "
-        '{"name": "f1", "arguments": {"a": 1}}'
-        '{"name": "f2", "arguments": {"b": 2}}'
-    )
     result = parser.parse_complete(body)
     assert result is not None
     assert result["content"] == "I will do two things."
@@ -588,30 +507,29 @@ def test_feed_then_parse_complete_same_instance():
     assert parser.has_tool_calls
 
 
-def test_pattern_matches_basic():
+@pytest.mark.parametrize(
-    assert _TOOL_CALL_HEAD_RE.search('{"name": "f"}')
+    "text,matches",
-
+    [
-
+        ('{"name": "f"}', True),
-def test_pattern_matches_with_whitespace():
+        ('{ "name" : "f"}', True),
-    assert _TOOL_CALL_HEAD_RE.search('{ "name" : "f"}')
+        ('{"other": 1}', False),
-
+        ('prefix {"name": "f", "args": {}}', True),
-
+        ('{"name": "f"}', True),  # match at start
-def test_pattern_no_match_without_name():
+        ('   {"name": "f"}', True),
-    assert _TOOL_CALL_HEAD_RE.search('{"other": 1}') is None
+    ],
-
+)
-
+def test_pattern_regex(text, matches):
-def test_pattern_match_mid_text():
+    result = _TOOL_CALL_HEAD_RE.search(text)
-    assert _TOOL_CALL_HEAD_RE.search('prefix {"name": "f", "args": {}}') is not None
+    if matches:
+        assert result is not None
+    else:
+        assert result is None
 
 
 def test_pattern_name_at_start():
     assert _TOOL_CALL_HEAD_RE.match('{"name": "f"}')
 
 
-def test_pattern_leading_whitespace():
-    assert _TOOL_CALL_HEAD_RE.search('   {"name": "f"}') is not None
-
-
 def test_factory_register_and_create():
     parser = ToolParserFactory.create("simple_json")
     assert isinstance(parser, BaseToolParser)
@@ -661,7 +579,6 @@ def test_feed_token_ids_do_not_affect_parsing():
         text, current_token_ids=[1, 2, 3], delta_token_ids=[3]
     )
     assert len(result_no) == len(result_with)
-    assert len(result_no) > 0
     assert (
         result_no[0]["tool_calls"][0]["function"]["name"]
         == result_with[0]["tool_calls"][0]["function"]["name"]

tests/module/test_encoder.py

@@ -1,6 +1,13 @@
+import json
+import os
+import tempfile
+
+import pytest
+import safetensors.torch as st
 import torch
 
 from astrai.config.model_config import EncoderConfig
+from astrai.model.automodel import AutoModel
 from astrai.model.encoder import EmbeddingEncoder
 
 TINY_CONFIG = dict(
@@ -14,92 +21,56 @@ TINY_CONFIG = dict(
     norm_eps=1e-5,
 )
 
+_device = "cuda" if torch.cuda.is_available() else "cpu"
 
-def test_encoder_forward_mean():
+
-    config = EncoderConfig(**TINY_CONFIG)
+def _make_model(**kwargs):
-    device = "cuda" if torch.cuda.is_available() else "cpu"
+    config = EncoderConfig(**{**TINY_CONFIG, **kwargs})
-    model = EmbeddingEncoder(config).to(device=device)
+    return EmbeddingEncoder(config).to(device=_device)
+
+
+@pytest.mark.parametrize("pooling_type", ["mean", "cls", "last"])
+def test_encoder_forward_pooling(pooling_type):
+    model = _make_model(pooling_type=pooling_type)
     model.eval()
 
     batch_size, seq_len = 2, 8
     input_ids = torch.randint(
-        0, config.vocab_size, (batch_size, seq_len), device=device
+        0, TINY_CONFIG["vocab_size"], (batch_size, seq_len), device=_device
     )
 
     with torch.no_grad():
         output = model(input_ids)
 
-    assert output.shape == (batch_size, config.dim)
+    assert output.shape == (batch_size, TINY_CONFIG["dim"])
-    assert not torch.isnan(output).any()
-
-
-def test_encoder_forward_cls():
-    config = EncoderConfig(**{**TINY_CONFIG, "pooling_type": "cls"})
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    model = EmbeddingEncoder(config).to(device=device)
-    model.eval()
-
-    batch_size, seq_len = 2, 8
-    input_ids = torch.randint(
-        0, config.vocab_size, (batch_size, seq_len), device=device
-    )
-
-    with torch.no_grad():
-        output = model(input_ids)
-
-    assert output.shape == (batch_size, config.dim)
-    assert not torch.isnan(output).any()
-
-
-def test_encoder_forward_last():
-    config = EncoderConfig(**{**TINY_CONFIG, "pooling_type": "last"})
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    model = EmbeddingEncoder(config).to(device=device)
-    model.eval()
-
-    batch_size, seq_len = 2, 8
-    input_ids = torch.randint(
-        0, config.vocab_size, (batch_size, seq_len), device=device
-    )
-
-    with torch.no_grad():
-        output = model(input_ids)
-
-    assert output.shape == (batch_size, config.dim)
     assert not torch.isnan(output).any()
 
 
 def test_encoder_forward_with_padding():
-    config = EncoderConfig(**TINY_CONFIG)
+    model = _make_model()
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    model = EmbeddingEncoder(config).to(device=device)
     model.eval()
 
     batch_size, seq_len = 2, 8
     input_ids = torch.randint(
-        0, config.vocab_size, (batch_size, seq_len), device=device
+        0, TINY_CONFIG["vocab_size"], (batch_size, seq_len), device=_device
     )
-    input_mask = torch.ones(batch_size, seq_len, dtype=torch.bool, device=device)
+    input_mask = torch.ones(batch_size, seq_len, dtype=torch.bool, device=_device)
     input_mask[:, 4:] = False
 
     with torch.no_grad():
         output = model(input_ids, input_mask=input_mask)
 
-    assert output.shape == (batch_size, config.dim)
+    assert output.shape == (batch_size, TINY_CONFIG["dim"])
     assert not torch.isnan(output).any()
 
 
 def test_encoder_normalize():
-    config = EncoderConfig(
+    model = _make_model(pooling_type="mean", normalize_embeddings=True)
-        **{**TINY_CONFIG, "pooling_type": "mean", "normalize_embeddings": True}
-    )
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    model = EmbeddingEncoder(config).to(device=device)
     model.eval()
 
     batch_size, seq_len = 2, 8
     input_ids = torch.randint(
-        0, config.vocab_size, (batch_size, seq_len), device=device
+        0, TINY_CONFIG["vocab_size"], (batch_size, seq_len), device=_device
     )
 
     with torch.no_grad():
@@ -110,24 +81,19 @@ def test_encoder_normalize():
 
 
 def test_encoder_register():
-    from astrai.model.automodel import AutoModel
-
     assert AutoModel.is_registered("embedding")
     cls = AutoModel.get_component_class("embedding")
     assert cls is EmbeddingEncoder
 
 
 def test_encoder_from_transformer_checkpoint():
-    config = EncoderConfig(**TINY_CONFIG)
+    model = _make_model()
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    model = EmbeddingEncoder(config).to(device=device)
-
     state_dict = model.state_dict()
     state_dict["lm_head.weight"] = torch.randn(
-        config.vocab_size, config.dim, device=device
+        TINY_CONFIG["vocab_size"], TINY_CONFIG["dim"], device=_device
     )
 
-    new_model = EmbeddingEncoder(config).to(device=device)
+    new_model = _make_model()
     new_model.load_state_dict(state_dict, strict=True)
 
     for key in model.state_dict():
@@ -135,12 +101,6 @@ def test_encoder_from_transformer_checkpoint():
 
 
 def test_encoder_save_load():
-    import json
-    import os
-    import tempfile
-
-    import safetensors.torch as st
-
     test_dir = tempfile.mkdtemp(prefix="encoder_test_")
     config_path = os.path.join(test_dir, "config.json")
     weights_path = os.path.join(test_dir, "model.safetensors")