77 lines
3.4 KiB
Python
77 lines
3.4 KiB
Python
from pathlib import Path
|
||
from typing import Literal
|
||
|
||
import yaml
|
||
from pydantic import Field, model_validator
|
||
|
||
from llm.config import ConfigModel, LLMConfig, GenerationConfig, CodexConfig
|
||
|
||
from .prompts import Prompts, load_prompts
|
||
|
||
DOMAINS = (
|
||
'manipulation', 'mechanical', 'fastening', 'press_fit', 'electrical',
|
||
'electronics', 'adhesive', 'welding', 'collaborative', 'inspection', 'operator_assistance',
|
||
)
|
||
|
||
|
||
class TeacherConfig(ConfigModel):
|
||
enabled: bool = False
|
||
fraction: float = Field(default=0.15, ge=0, le=1)
|
||
llm: LLMConfig | None = None
|
||
generation: GenerationConfig = Field(default_factory=GenerationConfig)
|
||
concurrency: int = Field(default=4, ge=1, le=128)
|
||
timeout_seconds: float = Field(default=180, gt=0)
|
||
retries: int = Field(default=2, ge=0, le=10)
|
||
|
||
@model_validator(mode='after')
|
||
def require_model(self):
|
||
if self.enabled and self.llm is None:
|
||
raise ValueError('teacher.llm обязателен при teacher.enabled: true')
|
||
return self
|
||
|
||
|
||
class Config(ConfigModel):
|
||
# Count means exported SFT rows, not episodes or API calls.
|
||
target_examples: int = Field(default=500_000, ge=0, strict=True)
|
||
seed: int = 20261021
|
||
output_dir: Path = Path('dataset/assembly')
|
||
modes: list[Literal['full', 'step', 'branch']] = Field(default_factory=lambda: ['full', 'step', 'branch'], min_length=1)
|
||
step_sampling: Literal['one', 'all'] = 'one'
|
||
shard_size: int = Field(default=1000, ge=1)
|
||
split: tuple[float, float, float] = (0.9, 0.05, 0.05)
|
||
domains: list[str] = Field(default_factory=lambda: list(DOMAINS), min_length=1)
|
||
min_components: int = Field(default=2, ge=1, le=12)
|
||
max_components: int = Field(default=6, ge=1, le=12)
|
||
blocked_fraction: float = Field(default=0.1, ge=0, le=1)
|
||
prompts_file: Path = Path("prompts.yaml")
|
||
prompts: Prompts = Field(default_factory=load_prompts)
|
||
teacher: TeacherConfig = Field(default_factory=TeacherConfig)
|
||
|
||
@model_validator(mode='after')
|
||
def check(self):
|
||
if len(set(self.modes)) != len(self.modes):
|
||
raise ValueError('modes не должны повторяться')
|
||
if len(set(self.domains)) != len(self.domains) or set(self.domains) - set(DOMAINS):
|
||
raise ValueError(f'Допустимые domains: {DOMAINS}')
|
||
if set(self.domains) - set(self.prompts.task_instructions):
|
||
raise ValueError("В task_instructions отсутствуют инструкции выбранных domains")
|
||
if self.min_components > self.max_components:
|
||
raise ValueError('min_components > max_components')
|
||
if any(v < 0 or v > 1 for v in self.split) or abs(sum(self.split) - 1) > 1e-9:
|
||
raise ValueError('split должен содержать три доли с суммой 1')
|
||
return self
|
||
|
||
|
||
def load_config(path: Path) -> Config:
|
||
path = path.resolve()
|
||
data = yaml.safe_load(path.read_text(encoding='utf-8')) or {}
|
||
prompts_path = (path.parent / data.get('prompts_file', 'prompts.yaml')).resolve()
|
||
data['prompts'] = load_prompts(prompts_path)
|
||
data['prompts_file'] = prompts_path
|
||
config = Config.model_validate(data)
|
||
config.output_dir = (path.parent / config.output_dir).resolve()
|
||
llm = config.teacher.llm
|
||
if isinstance(llm, CodexConfig) and llm.credential_path:
|
||
llm.credential_path = (path.parent / llm.credential_path).resolve()
|
||
return config
|