from transformers import PretrainedConfig

# ----------------------------
# Define Lumenspark Configuration
# ----------------------------

class LumensparkConfig(PretrainedConfig):
    """
    Configuration class for the Lumenspark model.
    Stores model hyperparameters like sequence length, embedding dimension, number of layers, and others.
    """
    model_type = "lumenspark"

    def __init__(
        self,
        seq_length=768,
        vocab_size=50257,
        embed_dim=768,
        depth=8,
        heads=12,
        dropout=1/17,
        k=384,
        rank=256,
        **kwargs
    ):
        super().__init__(**kwargs)
        self.vocab_size = vocab_size
        self.embed_dim = embed_dim
        self.depth = depth
        self.heads = heads
        self.seq_length = seq_length
        self.dropout = dropout
        self.k = k
        self.rank = rank

    def to_dict(self):
        """
        Converts the configuration parameters to a dictionary format.
        Useful for saving the configuration or inspecting model settings.
        """
        output = super().to_dict()
        output.update({
            "vocab_size": self.vocab_size,
            "embed_dim": self.embed_dim,
            "depth": self.depth,
            "heads": self.heads,
            "seq_length": self.seq_length,
            "dropout": self.dropout,
            "k": self.k,
            "rank": self.rank,
        })
        return output