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、配置序列化、工厂注册
This commit is contained in:
@@ -0,0 +1,88 @@
|
||||
"""Pipeline configuration for JSONL preprocessing."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Dict, Optional
|
||||
|
||||
|
||||
@dataclass
|
||||
class InputConfig:
|
||||
type: str = "chat"
|
||||
messages_key: str = "messages"
|
||||
prompt_key: str = "prompt"
|
||||
response_key: str = "response"
|
||||
text_key: str = "text"
|
||||
|
||||
|
||||
@dataclass
|
||||
class ProcessingConfig:
|
||||
max_seq_len: int = 2048
|
||||
min_chars: int = 50
|
||||
max_chars: int = 2_000_000
|
||||
deduplicate: bool = True
|
||||
max_items: Optional[int] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class OutputConfig:
|
||||
domain_key: Optional[str] = None
|
||||
storage_format: str = "bin"
|
||||
max_tokens_per_shard: int = 100_000_000
|
||||
|
||||
|
||||
@dataclass
|
||||
class PipelineConfig:
|
||||
version: int = 1
|
||||
input: InputConfig = field(default_factory=InputConfig)
|
||||
mask: Dict[str, str] = field(default_factory=dict)
|
||||
mask_default: str = "mask"
|
||||
preprocessing: ProcessingConfig = field(default_factory=ProcessingConfig)
|
||||
output: OutputConfig = field(default_factory=OutputConfig)
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
return {
|
||||
"version": self.version,
|
||||
"input": {
|
||||
"type": self.input.type,
|
||||
"messages_key": self.input.messages_key,
|
||||
"prompt_key": self.input.prompt_key,
|
||||
"response_key": self.input.response_key,
|
||||
"text_key": self.input.text_key,
|
||||
},
|
||||
"mask": self.mask,
|
||||
"mask_default": self.mask_default,
|
||||
"preprocessing": {
|
||||
"max_seq_len": self.preprocessing.max_seq_len,
|
||||
"min_chars": self.preprocessing.min_chars,
|
||||
"max_chars": self.preprocessing.max_chars,
|
||||
"deduplicate": self.preprocessing.deduplicate,
|
||||
"max_items": self.preprocessing.max_items,
|
||||
},
|
||||
"output": {
|
||||
"domain_key": self.output.domain_key,
|
||||
"storage_format": self.output.storage_format,
|
||||
"max_tokens_per_shard": self.output.max_tokens_per_shard,
|
||||
},
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: dict) -> PipelineConfig:
|
||||
return PipelineConfig(
|
||||
version=data.get("version", 1),
|
||||
input=InputConfig(**data.get("input", {})),
|
||||
mask=data.get("mask", {}),
|
||||
mask_default=data.get("mask_default", "mask"),
|
||||
preprocessing=ProcessingConfig(**data.get("preprocessing", {})),
|
||||
output=OutputConfig(**data.get("output", {})),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_json(cls, path: str) -> PipelineConfig:
|
||||
with open(path, "r", encoding="utf-8") as f:
|
||||
return cls.from_dict(json.load(f))
|
||||
|
||||
def to_json(self, path: str):
|
||||
with open(path, "w", encoding="utf-8") as f:
|
||||
json.dump(self.to_dict(), f, indent=2, ensure_ascii=False)
|
||||
Reference in New Issue
Block a user