adam/auto_training.py

"""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),
    )