from dataclasses import dataclass
from transformers import PretrainedConfig, AutoConfig


@dataclass
class NexusConfig:
    vocab_size: int = 50304
    max_seq_len: int = 512
    dim: int = 768
    num_layers: int = 10
    num_heads: int = 12
    num_kv_heads: int = 4
    multiple_of: int = 256
    ff_dim: int = 2048
    norm_eps: float = 1e-6
    dropout: float = 0.0
    
    batch_size: int = 4
    gradient_accumulation_steps: int = 8
    learning_rate: float = 3e-4
    min_lr: float = 3e-5
    weight_decay: float = 0.1
    beta1: float = 0.9
    beta2: float = 0.95
    warmup_steps: int = 500
    max_steps: int = 100000
    
    rope_theta: float = 500000.0
    
    use_flash_attention: bool = True
    
    data_path: str = "data"
    save_dir: str = "weights"
    
    seed: int = 42

nexus_config = NexusConfig()


class NexusSmAllConfig(PretrainedConfig):
    model_type = "nexus_small"

    def __init__(
        self,
        vocab_size=50304,
        max_seq_len=512,
        dim=768,
        num_layers=10,
        num_heads=12,
        num_kv_heads=4,
        multiple_of=256,
        ff_dim=2048,
        norm_eps=1e-6,
        rope_theta=500000.0,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.vocab_size = vocab_size
        self.max_seq_len = max_seq_len
        self.dim = dim
        self.num_layers = num_layers
        self.num_heads = num_heads
        self.num_kv_heads = num_kv_heads
        self.multiple_of = multiple_of
        self.ff_dim = ff_dim
        self.norm_eps = norm_eps
        self.rope_theta = rope_theta


AutoConfig.register("nexus_small", NexusSmAllConfig)