src/compliment_forest/training.py

from __future__ import annotations

from typing import Any, Literal, Protocol

from pydantic import BaseModel, ConfigDict, Field


class ChatTemplateTokenizer(Protocol):
    def apply_chat_template(
        self,
        messages: list[dict[str, str]],
        *,
        tokenize: bool,
        add_generation_prompt: bool,
        enable_thinking: bool,
    ) -> str: ...


class TextTrainingConfig(BaseModel):
    model_config = ConfigDict(extra="forbid", frozen=True)

    base_model: str = "openbmb/MiniCPM5-1B"
    dataset_id: str = "build-small-hackathon/compliment-forest-sft"
    model_id: str = "build-small-hackathon/compliment-forest-minicpm5-1b"
    adapter_id: str = "build-small-hackathon/compliment-forest-minicpm5-1b-lora"
    output_dir: str = "/training/compliment-forest-text"
    gguf_f16_name: str = "compliment-forest-minicpm5-1b.F16.gguf"
    gguf_q4_name: str = "compliment-forest-minicpm5-1b.Q4_K_M.gguf"
    max_length: int = Field(default=2048, ge=256)
    epochs: int = Field(default=2, ge=1)
    max_steps: int = -1
    train_limit: int | None = None
    validation_limit: int | None = None
    learning_rate: float = 2e-4
    per_device_batch_size: int = 2
    gradient_accumulation_steps: int = 8
    lora_rank: int = 16
    lora_alpha: int = 32
    lora_dropout: float = 0.05
    seed: int = 3407
    target_modules: tuple[str, ...] = (
        "q_proj",
        "k_proj",
        "v_proj",
        "o_proj",
        "gate_proj",
        "up_proj",
        "down_proj",
    )

    @classmethod
    def smoke(cls) -> TextTrainingConfig:
        return cls(
            epochs=1,
            max_steps=2,
            train_limit=16,
            validation_limit=8,
        )


class FluxTrainingConfig(BaseModel):
    model_config = ConfigDict(extra="forbid", frozen=True)

    base_model: str = "black-forest-labs/FLUX.1-dev"
    dataset_id: str = "build-small-hackathon/compliment-forest-watercolor"
    model_id: str = "build-small-hackathon/compliment-forest-flux-lora"
    trigger_token: str = "cmprst_forest"
    output_dir: str = "/training/compliment-forest-flux"
    resolution: int = 512
    max_train_steps: int = 500
    train_batch_size: int = 1
    gradient_accumulation_steps: int = 1
    learning_rate: float = 1e-4
    rank: int = 16
    lora_alpha: int = 16
    repeats: int = 4
    seed: int = 3407
    validation_prompt: str = (
        "a brave snail crossing a fern at dawn, cmprst_forest, soft watercolor "
        "storybook, gentle washes, warm dappled light, paper texture, kind expression, "
        "lots of negative space"
    )

    @classmethod
    def smoke(cls) -> FluxTrainingConfig:
        return cls(
            max_train_steps=2,
            repeats=1,
        )


TrainedForestStyle = Literal[
    "watercolor",
    "paper_cut",
    "moonlit_gouache",
    "botanical_ink",
]

_STYLE_TRAINING = {
    "watercolor": {
        "trigger": "cmprst_watercolor",
        "repo_suffix": "watercolor",
        "validation": "a gentle fox pausing beside ferns at dawn",
    },
    "paper_cut": {
        "trigger": "cmprst_papercut",
        "repo_suffix": "paper-cut",
        "validation": "a thoughtful badger beside layered woodland leaves",
    },
    "moonlit_gouache": {
        "trigger": "cmprst_moonlit",
        "repo_suffix": "moonlit-gouache",
        "validation": "a small owl resting in a moonlit pine clearing",
    },
    "botanical_ink": {
        "trigger": "cmprst_inkwash",
        "repo_suffix": "botanical-ink",
        "validation": "a patient hare beneath sparse woodland flowers",
    },
}


class FluxStyleTrainingConfig(BaseModel):
    model_config = ConfigDict(extra="forbid", frozen=True)

    style: TrainedForestStyle
    base_model: str = "black-forest-labs/FLUX.1-schnell"
    dataset_id: str = "thangvip/compliment-forest-multistyle-v2"
    dataset_config_name: str
    model_id: str
    trigger_token: str
    output_dir: str
    resolution: int = 512
    max_train_steps: int = 300
    train_batch_size: int = 1
    gradient_accumulation_steps: int = 1
    learning_rate: float = 1e-4
    rank: int = 16
    lora_alpha: int = 16
    repeats: int = 3
    seed: int = 3407
    guidance_scale: float = 0
    validation_prompt: str

    @classmethod
    def for_style(
        cls,
        style: TrainedForestStyle,
        *,
        smoke: bool = False,
    ) -> FluxStyleTrainingConfig:
        spec = _STYLE_TRAINING[style]
        config = cls(
            style=style,
            dataset_config_name=style,
            model_id=(
                f"thangvip/compliment-forest-{spec['repo_suffix']}-flux-lora-v2"
            ),
            trigger_token=spec["trigger"],
            output_dir=f"/training/compliment-forest-{spec['repo_suffix']}-flux",
            validation_prompt=f"{spec['trigger']}, {spec['validation']}",
        )
        if smoke:
            return config.model_copy(update={"max_train_steps": 2, "repeats": 1})
        return config


def format_training_example(
    example: dict[str, Any],
    tokenizer: ChatTemplateTokenizer,
) -> dict[str, str]:
    messages = example.get("messages")
    if not isinstance(messages, list):
        raise ValueError("training example must contain a messages list")
    return {
        "text": tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=False,
            enable_thinking=False,
        )
    }