refactor: migrate config system to Pydantic dataclasses

- Replace hand-rolled BaseConfig (from_dict/to_dict/_coerce/_unwrap_optional) with pydantic.dataclasses
- from_dict now uses cls(**d), to_dict uses dataclasses.asdict + json.dumps filter
- TrainConfig: required fields are now truly required (no default=None), delete manual validate()/__post_init__
- Remove dead required() helper and metadata={'help': ...} annotations
- Fix gradient_checkpointing_modules type from List[str] to List[type]
- Add pydantic>=2.0 as direct dependency in pyproject.toml
- Add numpy-style Parameters docstrings to all config classes
- Enable use_attribute_docstrings in BaseConfig for schema generation
- LoRAConfig also migrated to pydantic dataclass
This commit is contained in:
2026-07-30 08:25:32 +08:00
parent 6a97524db4
commit 164be9708b
7 changed files with 218 additions and 314 deletions
+35 -45
View File
@@ -5,9 +5,11 @@ modes, both driven declaratively through ``input.sections`` or
``input.sources``.
"""
from dataclasses import dataclass, field
from dataclasses import field
from typing import Dict, List, Optional
from pydantic.dataclasses import dataclass
from astrai.config.base import BaseConfig
@@ -25,6 +27,10 @@ class InputConfig(BaseConfig):
"chosen": {"sections": [{"field": "chosen", ...}]},
"rejected": {"sections": [{"field": "rejected", ...}]},
}}}
Args:
sections (Optional[List[Dict]]): Section list for single-output mode. Defaults to None.
sources (Optional[Dict[str, Dict]]): Source map for multi-output mode, DPO/GRPO. Defaults to None.
"""
sections: Optional[List[Dict]] = None
@@ -33,34 +39,17 @@ class InputConfig(BaseConfig):
@dataclass
class ProcessingConfig(BaseConfig):
"""Processing configuration.
"""Processing configuration for tokenization and packing.
Parameters
----------
max_seq_len : int
Maximum sequence length (default: 2048).
min_chars : int
Minimum number of characters to keep (default: 50).
max_chars : int
Maximum number of characters to keep (default: 2_000_000).
max_items : Optional[int]
Maximum number of items to process (default: None, unlimited).
batch_size : int
Number of records tokenized together (default: 256).
packing_strategy : str
How to pack sequences into a contiguous stream.
- ``"simple"``: sequential concatenation (default, backward compatible).
- ``"bfd"``: best-fit decreasing bin packing, minimises wasted tokens.
- ``"bfd_split"``: BFD with over-length sequences split into chunks.
max_packed_len : int
Maximum length of a packed bin. Sequences longer than this are
truncated or split depending on ``packing_strategy`` (default: 8192).
truncation_mode : str
How to truncate sequences longer than ``max_packed_len``.
- ``"keep_start"``: keep the first ``max_packed_len`` tokens (default).
- ``"keep_end"``: keep the last ``max_packed_len`` tokens.
Args:
max_seq_len (int): Maximum sequence length. Defaults to 2048.
min_chars (int): Minimum number of characters to keep. Defaults to 50.
max_chars (int): Maximum number of characters to keep. Defaults to 2_000_000.
max_items (Optional[int]): Maximum number of items to process, None=unlimited. Defaults to None.
batch_size (int): Number of records tokenized together. Defaults to 256.
packing_strategy (str): How to pack sequences: 'simple', 'bfd', or 'bfd_split'. Defaults to "simple".
max_packed_len (int): Maximum length of a packed bin. Defaults to 8192.
truncation_mode (str): How to truncate over-length sequences: 'keep_start' or 'keep_end'. Defaults to "keep_start".
"""
max_seq_len: int = 2048
@@ -75,24 +64,14 @@ class ProcessingConfig(BaseConfig):
@dataclass
class OutputConfig(BaseConfig):
"""Output configuration.
"""Output configuration for storage.
Parameters
----------
domain_key : Optional[str]
Domain key for the output store (default: None).
storage_format : str
Storage format, one of ``"bin"``, ``"jsonl"`` (default: ``"bin"``).
max_tokens_per_shard : int
Maximum tokens per shard before splitting (default: 100_000_000).
dtype : Dict[str, str]
Per-key dtype overrides, e.g. ``{"input_ids": "int32"}`` (default: {}).
position_ids_mode : Optional[str]
How to compute position_ids in packed sequences.
- ``"none"``: do not generate (default).
- ``"doc_reset"``: reset to 0 at each document boundary.
- ``"continuous"``: sequential 0, 1, 2, ... (pretrain, single doc).
Args:
domain_key (Optional[str]): Domain key for the output store. Defaults to None.
storage_format (str): Storage format: 'bin' or 'jsonl'. Defaults to "bin".
max_tokens_per_shard (int): Maximum tokens per shard before splitting. Defaults to 100_000_000.
dtype (Dict[str, str]): Per-key dtype overrides, e.g. {"input_ids": "int32"}. Defaults to {}.
position_ids_mode (str): Position ids mode: 'none', 'doc_reset', or 'continuous'. Defaults to "doc_reset".
"""
domain_key: Optional[str] = None
@@ -104,6 +83,17 @@ class OutputConfig(BaseConfig):
@dataclass
class PipelineConfig(BaseConfig):
"""Top-level preprocessing pipeline config.
Args:
version (int): Config schema version. Defaults to 1.
input (InputConfig): Input mapping config.
mask (Dict[str, str]): Per-field mask labels, e.g. {"system": "mask", "assistant": "train"}. Defaults to {}.
mask_default (str): Default mask label for unlisted fields. Defaults to "mask".
preprocessing (ProcessingConfig): Processing config.
output (OutputConfig): Output config.
"""
version: int = 1
input: InputConfig = field(default_factory=InputConfig)
mask: Dict[str, str] = field(default_factory=dict)