обновлен формат генерации датасета

assembly_dataset/config.py

@@ -0,0 +1,76 @@
+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