refactor : metric_logger 改用事件类型 (type=step/validation/epoch)

- 每种事件独立 schema，不再混入 null 字段
- 回调顺序 validation 移到 metric_logger 之前，确保 on_optimizer_step 先跑
- 用内部 _last_val_loss 代替 TrainContext.last_val_iter 判断新验证
- 修复 factory.py 未使用导入、evaluate_ifeval.py 多余 f 前缀

assets/docs/params.md

@@ -161,7 +161,7 @@ See [Inference Guide](inference.md) for HTTP API documentation.
 | `--top_k` | int | `30` | Top-k filtering |
 | `--top_p` | float | `0.95` | Nucleus sampling threshold |
 | `--batch_size` | int | `1` | Batch size for generation |
-| `--max_tokens` | int | `2048` | Maximum tokens to generate |
+| `--max_tokens` | int | model config `max_len` | Maximum tokens to generate |
 
 Usage:
 ```bash

astrai/config/model_config.py

@@ -20,6 +20,7 @@ class BaseModelConfig(BaseConfig):
     """Base config with ``model_type`` dispatch and file I/O."""
 
     model_type: Optional[str] = None
+    neftune_alpha: float = 0.0
 
 
 @dataclass
@@ -70,10 +71,12 @@ class EncoderConfig(BaseModelConfig):
     rope_theta: Optional[float] = None
     rope_scaling: Optional[dict] = None
 
+    attn_type: str = "gqa"
     n_heads: Optional[int] = None
     n_kv_heads: Optional[int] = None
     use_qk_norm: Optional[bool] = None
     use_gated_attention: Optional[bool] = None
 
+    ffn_type: str = "mlp"
     pooling_type: Optional[str] = None
     normalize_embeddings: Optional[bool] = None

astrai/factory.py

@@ -4,7 +4,6 @@ import inspect
 import sys
 from abc import ABC
 from typing import (
-    Any,
     Callable,
     Dict,
     ForwardRef,

astrai/model/components/attention.py

@@ -38,6 +38,7 @@ class GQA(nn.Module):
         norm_eps: float,
         use_gated_attention: bool,
         layer_id: int,
+        n_layers: int = 1,
     ):
         super().__init__()
         assert dim % n_heads == 0
@@ -55,7 +56,7 @@ class GQA(nn.Module):
         self.q_proj = Linear(dim, n_heads * self.head_dim)
         self.k_proj = Linear(dim, n_kv_heads * self.head_dim)
         self.v_proj = Linear(dim, n_kv_heads * self.head_dim)
-        self.o_proj = Linear(dim, dim)
+        self.o_proj = Linear(dim, dim, init_std=0.02 / (2 * n_layers) ** 0.5)
 
         if self.use_qk_norm:
             self.q_norm = RMSNorm(self.head_dim, norm_eps)
@@ -121,6 +122,7 @@ class MLA(nn.Module):
         use_qk_norm: bool,
         use_gated_attention: bool,
         layer_id: int,
+        n_layers: int = 1,
     ):
         super().__init__()
         self.dim = dim
@@ -148,7 +150,9 @@ class MLA(nn.Module):
             n_kv_heads * (2 * self.head_dim),
         )
 
-        self.o_proj = Linear(dim, dim, bias=False)
+        self.o_proj = Linear(
+            dim, dim, bias=False, init_std=0.02 / (2 * n_layers) ** 0.5
+        )
 
         if use_gated_attention:
             self.gate = Linear(dim, dim, bias=False)

astrai/model/components/decoder_block.py

@@ -1,3 +1,4 @@
+from dataclasses import asdict
 from typing import Optional
 
 import torch.nn as nn
@@ -10,35 +11,14 @@ from astrai.model.components.norm import RMSNorm
 
 
 class DecoderBlock(nn.Module):
-    def __init__(
-        self,
-        dim: int,
-        n_heads: int,
-        dim_ffn: int,
-        n_kv_heads: int,
-        norm_eps: float,
-        use_qk_norm: bool,
-        use_gated_attention: bool,
-        layer_id: int,
-        attn_type: str = "gqa",
-        ffn_type: str = "mlp",
-        **kwargs,
-    ):
+    def __init__(self, config, layer_id: int):
         super().__init__()
-        self.attention = AttnFactory.create(
-            attn_type,
-            dim=dim,
-            n_heads=n_heads,
-            n_kv_heads=n_kv_heads,
-            use_qk_norm=use_qk_norm,
-            norm_eps=norm_eps,
-            use_gated_attention=use_gated_attention,
-            layer_id=layer_id,
-            **kwargs,
-        )
-        self.input_norm = RMSNorm(dim, norm_eps)
-        self.post_attention_norm = RMSNorm(dim, norm_eps)
-        self.mlp = FFNFactory.create(ffn_type, dim, dim_ffn, **kwargs)
+        cfg = asdict(config)
+        cfg["down_init_std"] = 0.02 / (2 * config.n_layers) ** 0.5
+        self.attention = AttnFactory.create(config.attn_type, **cfg, layer_id=layer_id)
+        self.input_norm = RMSNorm(config.dim, config.norm_eps)
+        self.post_attention_norm = RMSNorm(config.dim, config.norm_eps)
+        self.mlp = FFNFactory.create(config.ffn_type, **cfg)
 
     def forward(
         self,

astrai/model/components/embedding.py

@@ -7,10 +7,13 @@ from torch import Tensor
 
 
 class Embedding(nn.Module):
-    def __init__(self, vocab_size: int, embedding_dim: int):
+    def __init__(self, vocab_size: int, embedding_dim: int, neftune_alpha: float = 0.0):
         super().__init__()
         self.weight = nn.Parameter(torch.empty((vocab_size, embedding_dim)))
-        self.neftune_noise_alpha = 0.0
+        self.neftune_noise_alpha = neftune_alpha
+
+    def set_neftune_alpha(self, alpha: float):
+        self.neftune_noise_alpha = alpha
 
     def reset_parameters(self):
         nn.init.normal_(self.weight, mean=0.0, std=0.02)

astrai/model/components/linear.py

@@ -5,13 +5,16 @@ from torch import Tensor
 
 
 class Linear(nn.Module):
-    def __init__(self, in_dim: int, out_dim: int, bias: bool = False):
+    def __init__(
+        self, in_dim: int, out_dim: int, bias: bool = False, init_std: float = 0.02
+    ):
         super().__init__()
         self.weight = nn.Parameter(torch.empty((out_dim, in_dim)))
         self.bias = nn.Parameter(torch.zeros(out_dim)) if bias else None
+        self.init_std = init_std
 
     def reset_parameters(self):
-        nn.init.kaiming_uniform_(self.weight, a=5**0.5)
+        nn.init.normal_(self.weight, mean=0.0, std=self.init_std)
         if self.bias is not None:
             fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight)
             bound = 1 / (fan_in**0.5)

astrai/model/components/mlp.py

@@ -13,11 +13,11 @@ class FFNFactory(BaseFactory[nn.Module]):
 
 @FFNFactory.register("mlp")
 class MLP(nn.Module):
-    def __init__(self, dim: int, dim_ffn: int):
+    def __init__(self, dim: int, dim_ffn: int, down_init_std: float = 0.02):
         super().__init__()
         self.up = Linear(dim, dim_ffn)
         self.gate = Linear(dim, dim_ffn)
-        self.down = Linear(dim_ffn, dim)
+        self.down = Linear(dim_ffn, dim, init_std=down_init_std)
 
     def forward(self, x: Tensor) -> Tensor:
         gated = self.up(x) * F.silu(self.gate(x))
@@ -35,6 +35,7 @@ class DeepSeekMoE(nn.Module):
         n_shared_experts: int = 1,
         n_activated_experts: int = 2,
         topk_method: str = "greedy",
+        n_layers: int = 1,
     ):
         super().__init__()
         self.dim = dim
@@ -44,12 +45,20 @@ class DeepSeekMoE(nn.Module):
         self.topk_method = topk_method
 
         self.router = Linear(dim, n_routed_experts, bias=False)
+        moe_scale = 1 / max(n_shared_experts, 1) + 1 / n_activated_experts
+        down_init_std = 0.02 / (2 * n_layers * moe_scale) ** 0.5
 
         self.shared_experts = nn.ModuleList(
-            [MLP(dim, dim_ffn) for _ in range(n_shared_experts)]
+            [
+                MLP(dim, dim_ffn, down_init_std=down_init_std)
+                for _ in range(n_shared_experts)
+            ]
         )
         self.routed_experts = nn.ModuleList(
-            [MLP(dim, dim_ffn) for _ in range(n_routed_experts)]
+            [
+                MLP(dim, dim_ffn, down_init_std=down_init_std)
+                for _ in range(n_routed_experts)
+            ]
         )
 
     def forward(self, x: Tensor) -> Tensor:

astrai/model/encoder.py

@@ -23,22 +23,12 @@ class EmbeddingEncoder(AutoModel):
         self.rotary_embedding = RotaryEmbedding(
             rope_dim, config.max_len, rope_base, rope_scaling=config.rope_scaling
         )
-        self.embed_tokens = Embedding(config.vocab_size, config.dim)
+        self.embed_tokens = Embedding(
+            config.vocab_size, config.dim, neftune_alpha=config.neftune_alpha
+        )
 
         self.layers = nn.ModuleList(
-            [
-                DecoderBlock(
-                    config.dim,
-                    config.n_heads,
-                    config.dim_ffn,
-                    config.n_kv_heads,
-                    config.norm_eps,
-                    config.use_qk_norm,
-                    config.use_gated_attention,
-                    layer_id,
-                )
-                for layer_id in range(config.n_layers)
-            ]
+            [DecoderBlock(config, layer_id) for layer_id in range(config.n_layers)]
         )
 
         self.norm = RMSNorm(config.dim, config.norm_eps)

astrai/model/transformer.py

@@ -59,31 +59,12 @@ class AutoRegressiveLM(AutoModel):
         self.rotary_embedding = RotaryEmbedding(
             rope_dim, config.max_len, rope_base, rope_scaling=config.rope_scaling
         )
-        self.embed_tokens = Embedding(config.vocab_size, config.dim)
+        self.embed_tokens = Embedding(
+            config.vocab_size, config.dim, neftune_alpha=config.neftune_alpha
+        )
 
         self.layers = nn.ModuleList(
-            [
-                DecoderBlock(
-                    config.dim,
-                    config.n_heads,
-                    config.dim_ffn,
-                    config.n_kv_heads,
-                    config.norm_eps,
-                    config.use_qk_norm,
-                    config.use_gated_attention,
-                    layer_id,
-                    attn_type=config.attn_type,
-                    ffn_type=config.ffn_type,
-                    n_routed_experts=config.n_routed_experts,
-                    n_shared_experts=config.n_shared_experts,
-                    n_activated_experts=config.n_activated_experts,
-                    topk_method=config.topk_method,
-                    kv_lora_rank=config.kv_lora_rank,
-                    qk_nope_head_dim=config.qk_nope_head_dim,
-                    qk_rope_head_dim=config.qk_rope_head_dim,
-                )
-                for layer_id in range(config.n_layers)
-            ]
+            [DecoderBlock(config, layer_id) for layer_id in range(config.n_layers)]
         )
 
         self.norm = RMSNorm(config.dim, config.norm_eps)

astrai/trainer/train_callback.py

@@ -237,6 +237,7 @@ class MetricLoggerCallback(TrainCallback):
         metrics: List[str] = None,
     ):
         self.last_log_iter = 0
+        self._last_val_loss = None
         self.save_interval = save_interval
         self.log_interval = log_interval
         self.metrics = metrics or ["loss", "lr"]
@@ -258,46 +259,54 @@ class MetricLoggerCallback(TrainCallback):
             "grad_nan_num": ctx_get_grad_nan_num,
         }
 
-    def _get_log_data(self, context: TrainContext):
-        data = {
+    def _metrics(self, context: TrainContext, names):
+        return {
+            m: self._metric_funcs[m](context)
+            for m in names
+            if self._metric_funcs[m](context) is not None
+        }
+
+    @only_on_rank(0)
+    def _append(self, event_type: str, context: TrainContext, **extra):
+        entry = {
+            "type": event_type,
             "timestamp": time.strftime("%Y-%m-%dT%H:%M:%S"),
             "epoch": context.epoch,
             "iter": context.iteration,
+            **extra,
         }
-        for m in self.metrics:
-            val = self._metric_funcs[m](context)
-            if val is not None:
-                data[m] = val
-        return data
+        self.log_cache.append(entry)
 
     @only_on_rank(0)
-    def _add_log(self, log_data):
-        self.log_cache.append(log_data)
-
-    @only_on_rank(0)
-    def _save_log(self, epoch, iter):
+    def _flush(self, epoch, iter):
         log_file = self.log_dir / f"epoch_{epoch}_iter_{iter}_metric.jsonl"
         log_file.parent.mkdir(parents=True, exist_ok=True)
-
         with open(log_file, "w") as f:
             for log in self.log_cache:
                 f.write(json.dumps(log) + "\n")
 
     def on_batch_end(self, context):
         if context.iteration % self.log_interval == 0:
-            log_data = self._get_log_data(context)
-            self._add_log(log_data)
-
+            step_metrics = [m for m in self.metrics if m != "val_loss"]
+            self._append("step", context, **self._metrics(context, step_metrics))
         if context.iteration - self.last_log_iter >= self.save_interval:
-            self._save_log(context.epoch, context.iteration)
+            self._flush(context.epoch, context.iteration)
             self.last_log_iter = context.iteration
 
+    def on_optimizer_step(self, context):
+        if context.val_loss is not None and context.val_loss != self._last_val_loss:
+            self._append("validation", context, val_loss=context.val_loss)
+            self._last_val_loss = context.val_loss
+
+    def on_epoch_end(self, context):
+        self._append("epoch", context)
+
     def on_train_end(self, context):
         if context.iteration != self.last_log_iter:
-            self._save_log(context.epoch, context.iteration)
+            self._flush(context.epoch, context.iteration)
 
     def on_error(self, context):
-        self._save_log(context.epoch, context.iteration)
+        self._flush(context.epoch, context.iteration)
 
 
 @CallbackFactory.register("validation")

astrai/trainer/train_context.py

@@ -63,7 +63,6 @@ class TrainContextBuilder:
 
         model = cfg.model_fn()
         model = model.to(device=device)
-        model.embed_tokens.neftune_noise_alpha = cfg.neftune_alpha
 
         model_config = {}
         if self._resume_dir:

astrai/trainer/trainer.py

@@ -34,6 +34,7 @@ class Trainer:
                 cfg.ckpt_dir,
                 cfg.ckpt_interval,
             ),
+            CallbackFactory.create("validation"),
             CallbackFactory.create(
                 "metric_logger",
                 log_dir=cfg.log_dir,
@@ -43,7 +44,6 @@ class Trainer:
             ),
             CallbackFactory.create("progress_bar", cfg.n_epoch),
             CallbackFactory.create("gradient_clipping", cfg.max_grad_norm),
-            CallbackFactory.create("validation"),
         ]
         return callbacks
 

scripts/eval/evaluate_ifeval.py

@@ -571,7 +571,7 @@ def main():
     print(f"  Unsupported:    {summary['unsupported_constraints']}")
     print(f"{'=' * 60}")
 
-    print(f"\nPer-type accuracy:")
+    print("\nPer-type accuracy:")
     for inst_id, stats in sorted(summary["per_type_accuracy"].items()):
         print(
             f"  {inst_id:50s}  {stats['accuracy']:.2%}  "

scripts/eval/evaluate_ppl.py

@@ -86,7 +86,7 @@ def process_file(
 
 
 if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="Run perplexity with a Khaosz model.")
+    parser = argparse.ArgumentParser(description="Perplexity evaluation on JSONL text.")
     parser.add_argument(
         "--param_path", type=str, required=True, help="Path to the model directory."
     )

scripts/tools/generate.py

@@ -1,5 +1,6 @@
 import argparse
 import json
+from typing import Optional
 
 import torch
 
@@ -17,7 +18,7 @@ def processor(
     top_p: float,
     question_key: str,
     response_key: str,
-    max_tokens: int,
+    max_tokens: Optional[int],
     batch_size: int,
 ):
     # Load model and tokenizer
@@ -72,7 +73,7 @@ def processor(
 
 
 if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="Run generate with a Khaosz model.")
+    parser = argparse.ArgumentParser(description="Batch generation from JSONL file.")
 
     parser.add_argument(
         "--param_path", type=str, required=True, help="Path to the model directory."
@@ -93,36 +94,42 @@ if __name__ == "__main__":
         "--question_key",
         type=str,
         default="question",
-        help="Key for the question in the input JSON.",
+        help="Key for the question in the input JSON (default: question).",
     )
     parser.add_argument(
         "--response_key",
         type=str,
         default="response",
-        help="Key for the response in the output JSON.",
+        help="Key for the response in the output JSON (default: response).",
     )
     parser.add_argument(
         "--temperature",
         type=float,
         default=0.60,
-        help="Temperature for generating responses.",
+        help="Temperature for generating responses (default: 0.60).",
     )
     parser.add_argument(
-        "--top_k", type=int, default=30, help="Top-k value for generating responses."
+        "--top_k",
+        type=int,
+        default=30,
+        help="Top-k value for generating responses (default: 30).",
     )
     parser.add_argument(
         "--top_p",
         type=float,
         default=0.95,
-        help="Top-p value for generating responses.",
+        help="Top-p value for generating responses (default: 0.95).",
     )
     parser.add_argument(
-        "--batch_size", type=int, default=1, help="Batch size for generating responses."
+        "--batch_size",
+        type=int,
+        default=1,
+        help="Batch size for generating responses (default: 1).",
     )
     parser.add_argument(
         "--max_tokens",
         type=int,
-        default=2048,
+        default=None,
         help="Maximum tokens to generate (default: model config max_len).",
     )
 

scripts/tools/train.py

@@ -221,6 +221,44 @@ def parse_args() -> argparse.Namespace:
         help="NEFTune noise alpha (0=disabled, typical: 5.0).",
     )
 
+    parser.add_argument(
+        "--schedule_type",
+        type=str,
+        default="cosine",
+        choices=["cosine", "sgdr", "wsd"],
+        help="Learning rate scheduler type.",
+    )
+    parser.add_argument(
+        "--min_rate",
+        type=float,
+        default=None,
+        help="Minimum LR as fraction of base LR. Uses scheduler default if not set (cosine/sgdr: 0.05, wsd: 0.0).",
+    )
+    parser.add_argument(
+        "--cycle_length",
+        type=int,
+        default=None,
+        help="SGDR first cycle length in steps. Defaults to total_steps - warmup_steps.",
+    )
+    parser.add_argument(
+        "--t_mult",
+        type=int,
+        default=2,
+        help="SGDR cycle length multiplier per restart.",
+    )
+    parser.add_argument(
+        "--stable_steps",
+        type=int,
+        default=None,
+        help="WSD stable plateau steps. Required when --schedule_type wsd.",
+    )
+    parser.add_argument(
+        "--decay_steps",
+        type=int,
+        default=None,
+        help="WSD decay steps. Defaults to total_steps - warmup_steps - stable_steps.",
+    )
+
     args = parser.parse_args()
 
     return args
@@ -300,6 +338,12 @@ def train(
     master_port: str,
     start_method: str,
     neftune_alpha: float,
+    schedule_type: str,
+    min_rate: float,
+    cycle_length: int,
+    t_mult: int,
+    stable_steps: int,
+    decay_steps: int,
 ):
     assert train_type in ["seq", "sft", "dpo", "grpo"]
     assert os.path.exists(param_path)
@@ -309,6 +353,7 @@ def train(
     # Load config
     config_path = os.path.join(param_path, "config.json")
     config = AutoRegressiveLMConfig.from_file(config_path)
+    config.neftune_alpha = neftune_alpha
 
     if window_size is None:
         window_size = config.max_len
@@ -348,14 +393,30 @@ def train(
         len(dataset), n_epoch, batch_per_device, nprocs, grad_accum_steps
     )
     warmup_steps = int(warmup_ratio * total_steps)
+    warmup_steps = min(warmup_steps, total_steps)
+
+    scheduler_kwargs = {"warmup_steps": warmup_steps}
+
+    if schedule_type == "cosine":
+        scheduler_kwargs["lr_decay_steps"] = total_steps - warmup_steps
+    elif schedule_type == "sgdr":
+        scheduler_kwargs["cycle_length"] = cycle_length or (total_steps - warmup_steps)
+        scheduler_kwargs["t_mult"] = t_mult
+    elif schedule_type == "wsd":
+        remaining = total_steps - warmup_steps
+        stable_steps_ = stable_steps or max(1, int(remaining * 0.8))
+        scheduler_kwargs["stable_steps"] = stable_steps_
+        scheduler_kwargs["decay_steps"] = max(
+            1, decay_steps or (remaining - stable_steps_)
+        )
+
+    if min_rate is not None:
+        scheduler_kwargs["min_rate"] = min_rate
 
     scheduler_fn = partial(
         create_scheduler,
-        **{
-            "schedule_type": "cosine",
-            "warmup_steps": min(warmup_steps, total_steps),
-            "lr_decay_steps": total_steps - min(warmup_steps, total_steps),
-        },
+        schedule_type=schedule_type,
+        **scheduler_kwargs,
     )
 
     grad_ckpt_modules = [DecoderBlock] if gradient_checkpointing else []

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(

tests/inference/test_tool_parser.py

@@ -13,67 +13,28 @@ from astrai.inference.api.tool_parser import (
 )
 
 
-def test_scan_complete_simple():
-    end, complete = _scan_json('{"key": "value"}', 0)
-    assert complete is True
-    assert end == len('{"key": "value"}')
-
-
-def test_scan_complete_nested():
-    text = '{"outer": {"inner": 1}}'
+@pytest.mark.parametrize(
+    "text,expected_complete,check_end_eq_len",
+    [
+        ('{"key": "value"}', True, True),
+        ('{"outer": {"inner": 1}}', True, True),
+        ('{"key": "value"', False, False),
+        ('{"outer": {"inner": 1}', False, False),
+        ('{"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 = (
-        '{"name": "batch", '
-        '"arguments": {"rows": [{"id": 1, "val": "a"}, {"id": 2, "val": "b"}]}}'
-    )
+    text = '{"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():
-    result = _find_partial_tool_call('{"name": "func", "arguments": {"city"')
+@pytest.mark.parametrize(
+    "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["complete"] is False
-
-
-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"
+        assert result["name"] == expected_name
+        if expected_complete is not None:
+            assert result["complete"] is expected_complete
 
 
 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 = (
-        '{"name": "get_weather", "arguments": {"city": "Beijing"}}'
-        '{"name": "get_time", "arguments": {"tz": "Asia/Shanghai"}}'
-    )
+    body = '{"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", '
-        '"arguments": {'
-        '"to": ["a@b.com", "c@d.com"], '
-        '"cc": null, '
-        '"subject": "Hello World", '
-        '"body": "This is a test email.", '
-        '"priority": 1, '
-        '"attachments": false'
-        "}}"
+        '{"name": "send_email", "arguments": {'
+        '"to": ["a@b.com", "c@d.com"], "cc": null, '
+        '"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 = (
-        "I will do two things. "
-        '{"name": "f1", "arguments": {"a": 1}}'
-        '{"name": "f2", "arguments": {"b": 2}}'
-    )
+    body = '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():
-    assert _TOOL_CALL_HEAD_RE.search('{"name": "f"}')
-
-
-def test_pattern_matches_with_whitespace():
-    assert _TOOL_CALL_HEAD_RE.search('{ "name" : "f"}')
-
-
-def test_pattern_no_match_without_name():
-    assert _TOOL_CALL_HEAD_RE.search('{"other": 1}') is None
-
-
-def test_pattern_match_mid_text():
-    assert _TOOL_CALL_HEAD_RE.search('prefix {"name": "f", "args": {}}') is not None
+@pytest.mark.parametrize(
+    "text,matches",
+    [
+        ('{"name": "f"}', True),
+        ('{ "name" : "f"}', True),
+        ('{"other": 1}', False),
+        ('prefix {"name": "f", "args": {}}', True),
+        ('{"name": "f"}', True),  # match at start
+        ('   {"name": "f"}', True),
+    ],
+)
+def test_pattern_regex(text, matches):
+    result = _TOOL_CALL_HEAD_RE.search(text)
+    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)
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    model = EmbeddingEncoder(config).to(device=device)
+
+def _make_model(**kwargs):
+    config = EncoderConfig(**{**TINY_CONFIG, **kwargs})
+    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 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 output.shape == (batch_size, TINY_CONFIG["dim"])
     assert not torch.isnan(output).any()
 
 
 def test_encoder_forward_with_padding():
-    config = EncoderConfig(**TINY_CONFIG)
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    model = EmbeddingEncoder(config).to(device=device)
+    model = _make_model()
     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(
-        **{**TINY_CONFIG, "pooling_type": "mean", "normalize_embeddings": True}
-    )
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    model = EmbeddingEncoder(config).to(device=device)
+    model = _make_model(pooling_type="mean", normalize_embeddings=True)
     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)
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    model = EmbeddingEncoder(config).to(device=device)
-
+    model = _make_model()
     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")