from pydantic import Field

from ..base import BaseModel


class LoraConfig(BaseModel):
    base_model: str
    rank: int = Field(ge=4, default=32, le=64)
    seed: int | None = None
    train_mlp: bool = True
    train_attn: bool = True
    train_unembed: bool = True
