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:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user