from dataclasses import dataclass, field
from typing import Optional

from simple_parsing.helpers import Serializable


@dataclass
class LoraArgs(Serializable):
    enable: bool = True
    rank: int = 16
    dropout: float = 0.0
    scaling: float = 2.0

    def __post_init__(self):
        if self.enable:
            assert self.rank > 0
            assert self.scaling > 0.0


@dataclass
class MoeArgs(Serializable):
    num_experts: int = 8
    num_experts_per_tok: int = 2


@dataclass
class ModelArgs(Serializable):
    dim: int
    n_layers: int
    head_dim: int
    hidden_dim: int
    n_heads: int
    n_kv_heads: int
    norm_eps: float
    vocab_size: int
    rope_theta: float = 10000.0

    lora: LoraArgs = field(default_factory=LoraArgs)
    moe: Optional[MoeArgs] = None