"""Intent-level AUTO policies for ADAM training requests. Natural-language parsing belongs in the planner. This module deliberately does not inspect prompt wording: it turns an already-selected trainer, profile, and dataset size into transparent, reproducible training settings. """ from __future__ import annotations from dataclasses import dataclass from enum import Enum from typing import Any from adam.model_profiles import ModelProfile from adam.models import SystemSnapshot from adam.recommendations import SettingsRecommendation, recommend_for_profile class TrainingProfile(str, Enum): TEST = "test" BALANCED = "balanced" QUALITY = "quality" OVERNIGHT = "overnight" @dataclass(frozen=True, slots=True) class AutoTrainingPlan: """Resolved, explainable settings for one training run.""" profile: TrainingProfile dataset_target: int epochs: int settings: dict[str, Any] summary: str reasons: tuple[str, ...] = () warnings: tuple[str, ...] = () def profile_from_request(request: str) -> TrainingProfile: """Map only stable, user-facing intent modifiers to a named policy.""" lowered = request.casefold() if any(word in lowered for word in ("overnight", "all night", "long run")): return TrainingProfile.OVERNIGHT if any(phrase in lowered for phrase in ("high quality", "best quality", "really good", "quality")): return TrainingProfile.QUALITY if any(word in lowered for word in ("quick", "quickly", "test", "small", "smoke test")): return TrainingProfile.TEST return TrainingProfile.BALANCED def dataset_target_for(trainer: str, profile: TrainingProfile) -> int: """Choose a collection target, never a random count. These are conservative collection targets. The later recommendation is calculated from the actual usable count when an existing dataset is known. """ targets = { "lora": {TrainingProfile.TEST: 40, TrainingProfile.BALANCED: 150, TrainingProfile.QUALITY: 300, TrainingProfile.OVERNIGHT: 500}, "ddpm": {TrainingProfile.TEST: 100, TrainingProfile.BALANCED: 400, TrainingProfile.QUALITY: 800, TrainingProfile.OVERNIGHT: 1_200}, "flow": {TrainingProfile.TEST: 100, TrainingProfile.BALANCED: 400, TrainingProfile.QUALITY: 800, TrainingProfile.OVERNIGHT: 1_200}, "inrflow": {TrainingProfile.TEST: 80, TrainingProfile.BALANCED: 300, TrainingProfile.QUALITY: 600, TrainingProfile.OVERNIGHT: 900}, } return targets.get(trainer, targets["ddpm"])[profile] def resolve_auto_training( profile: ModelProfile, *, trainer: str, policy: TrainingProfile, dataset_items: int, snapshot: SystemSnapshot | None = None, ) -> AutoTrainingPlan: """Resolve a named policy through ADAM's existing exposure-aware recommender.""" recommendation: SettingsRecommendation = recommend_for_profile( profile, dataset_items=max(10, dataset_items), snapshot=snapshot, ) epoch_multiplier = { TrainingProfile.TEST: 0.20, TrainingProfile.BALANCED: 1.00, TrainingProfile.QUALITY: 1.35, TrainingProfile.OVERNIGHT: 1.80, }[policy] # Preserve safe bounds from the recommendation; profiles express their own # architecture-specific baseline rather than sharing a global epoch range. minimum = 3 if policy is TrainingProfile.TEST else 10 maximum = 1_000 if policy is TrainingProfile.OVERNIGHT else 600 epochs = max(minimum, min(maximum, round(recommendation.epochs * epoch_multiplier))) settings = dict(recommendation.settings) if "save_every" in settings: settings["save_every"] = max(1, min(int(settings["save_every"]), max(1, epochs // 4))) if "preview_every" in settings: settings["preview_every"] = max(1, min(int(settings["preview_every"]), max(1, epochs // 5))) target = dataset_target_for(trainer, policy) summary = ( f"{policy.value.title()} AUTO policy: target {target:,} source images; " f"{epochs:,} epoch budget based on {max(10, dataset_items):,} expected usable items." ) return AutoTrainingPlan( profile=policy, dataset_target=target, epochs=epochs, settings=settings, summary=summary, reasons=tuple(recommendation.reasons), warnings=tuple(recommendation.warnings), )