"""Hugging Face configuration for the educational FineWeb decoder."""

from transformers import PreTrainedConfig


class FineWebConfig(PreTrainedConfig):
    model_type = "fineweb_decoder"

    def __init__(
        self,
        vocab_size=16384,
        n_layers=8,
        d_model=256,
        n_heads=4,
        mlp_hidden=960,
        context_length=1024,
        rms_norm_eps=1e-6,
        **kwargs,
    ):
        kwargs.setdefault("tie_word_embeddings", True)
        kwargs.setdefault("use_cache", False)
        super().__init__(**kwargs)
        self.vocab_size = vocab_size
        self.n_layers = n_layers
        self.d_model = d_model
        self.n_heads = n_heads
        self.mlp_hidden = mlp_hidden
        self.context_length = context_length
        self.max_position_embeddings = context_length
        self.rms_norm_eps = rms_norm_eps
        self.hidden_size = d_model
        self.num_hidden_layers = n_layers
        self.num_attention_heads = n_heads