File size: 4,371 Bytes
f8c73f9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 | """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),
)
|