refactor : 基于声明式 JSON 配置的预处理管线重构

- 用工厂注册的 MaskBuilder（chat/instruction/text）替换硬编码的 _transform_* 方法 - mask 规则以 role-to-action 映射声明在配置中，与 chat_template 完全解耦 - 单次编码 + role-span 追踪替代两次编码 + 长度差计算 mask 的方式 - 支持多轮对话训练：所有 assistant 轮次参与训练，而非仅最后一轮 - 新建 astrai.preprocessing 包（builder.py + pipeline.py），删除 astrai/preprocess.py - CLI 精简为 --config 参数，所有参数通过 PipelineConfig JSON 配置 - 新增 PipelineConfig、InputConfig、ProcessingConfig、OutputConfig dataclass - 文档：assets/docs/preprocessing.md - 27 个测试覆盖 mask builder、pipeline、配置序列化、工厂注册
2026-05-30 20:45:09 +08:00
parent 138c5bcc08
commit 69207e2c57
10 changed files with 1170 additions and 386 deletions
@@ -0,0 +1,161 @@
+"""Mask building strategies for preprocessing pipeline.
+
+Each builder knows how to tokenize one input format and construct
+the loss_mask according to declarative mask rules from the config.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import List, Optional
+
+from astrai.factory import BaseFactory
+
+
+class BaseMaskBuilder(ABC):
+    """Convert a JSONL item into token ids and optional loss_mask."""
+
+    @abstractmethod
+    def build(self, item: dict, config, tokenizer) -> Optional[dict]:
+        """Build ``{ids, loss_mask?, domain}`` from a JSONL record.
+
+        Returns ``None`` to skip the item entirely.
+        """
+        ...
+
+
+class MaskBuilderFactory(BaseFactory["BaseMaskBuilder"]):
+    @classmethod
+    def _validate_component(cls, component_cls: type):
+        if not issubclass(component_cls, BaseMaskBuilder):
+            raise TypeError(
+                f"{component_cls.__name__} must inherit from BaseMaskBuilder"
+            )
+
+
+def _extract_domain(item: dict, domain_key: Optional[str]) -> str:
+    if not domain_key:
+        return "__default__"
+    val = item.get(domain_key, "__default__")
+    return val if isinstance(val, str) else "__default__"
+
+
+@MaskBuilderFactory.register("chat")
+class ChatMaskBuilder(BaseMaskBuilder):
+    """Mask by role via message-level tokenisation with role-span tracking.
+
+    For each message, renders the chat template for that single message,
+    encodes individually, and records its token span + role action.
+    The concatenated sequence receives a loss_mask built from span rules.
+    """
+
+    def build(self, item: dict, config, tokenizer) -> Optional[dict]:
+        messages = item.get(config.input.messages_key)
+        if not isinstance(messages, list) or not messages:
+            return None
+
+        all_ids: List[int] = []
+        spans: List[tuple] = []
+
+        if tokenizer.bos_token_id is not None:
+            all_ids.append(tokenizer.bos_token_id)
+
+        for msg in messages:
+            role = msg.get("role", "")
+            action = config.mask.get(role, config.mask_default)
+
+            rendered = tokenizer.apply_chat_template(
+                [msg], tokenize=False, add_generation_prompt=False
+            )
+            ids = tokenizer.encode(rendered, add_special_tokens=False)
+
+            start = len(all_ids)
+            all_ids.extend(ids)
+            spans.append((start, len(all_ids), action))
+
+        if len(all_ids) <= 1:
+            return None
+
+        max_len = config.preprocessing.max_seq_len
+        all_ids = all_ids[:max_len]
+
+        loss_mask = [0] * len(all_ids)
+        for start, end, action in spans:
+            if start >= len(all_ids):
+                break
+            e = min(end, len(all_ids))
+            if action == "train":
+                loss_mask[start:e] = [1] * (e - start)
+
+        return {
+            "ids": all_ids,
+            "loss_mask": loss_mask,
+            "domain": _extract_domain(item, config.output.domain_key),
+        }
+
+
+@MaskBuilderFactory.register("instruction")
+class InstructionMaskBuilder(BaseMaskBuilder):
+    """Mask by prompt / response field boundary.
+
+    Encodes prompt and response independently, then fills mask
+    according to ``prompt`` / ``response`` entries in the mask config.
+    """
+
+    def build(self, item: dict, config, tokenizer) -> Optional[dict]:
+        prompt = str(item.get(config.input.prompt_key, ""))
+        response = str(item.get(config.input.response_key, ""))
+
+        if not prompt.strip() and not response.strip():
+            return None
+
+        prompt_ids = tokenizer.encode(prompt, add_special_tokens=True)
+        response_ids = tokenizer.encode(response, add_special_tokens=False)
+
+        max_len = config.preprocessing.max_seq_len
+        full_ids = (prompt_ids + response_ids)[:max_len]
+
+        prompt_action = config.mask.get("prompt", config.mask_default)
+        response_action = config.mask.get("response", config.mask_default)
+
+        p_len = min(len(prompt_ids), len(full_ids))
+        r_len = len(full_ids) - p_len
+
+        loss_mask = []
+        if prompt_action == "train":
+            loss_mask += [1] * p_len
+        else:
+            loss_mask += [0] * p_len
+
+        if response_action == "train":
+            loss_mask += [1] * r_len
+        else:
+            loss_mask += [0] * r_len
+
+        return {
+            "ids": full_ids,
+            "loss_mask": loss_mask,
+            "domain": _extract_domain(item, config.output.domain_key),
+        }
+
+
+@MaskBuilderFactory.register("text")
+class TextMaskBuilder(BaseMaskBuilder):
+    """Plain tokenisation — no mask, used for pre-training data."""
+
+    def build(self, item: dict, config, tokenizer) -> Optional[dict]:
+        text = item.get(config.input.text_key, "")
+        if not isinstance(text, str) or not text.strip():
+            return None
+
+        pp = config.preprocessing
+        if not (pp.min_chars <= len(text) <= pp.max_chars):
+            return None
+
+        ids = tokenizer.encode(text, add_special_tokens=True)
+        ids = ids[: pp.max_seq_len]
+
+        return {
+            "ids": ids,
+            "domain": _extract_domain(item, config.output.domain_key),
+        }