from transformers import PretrainedConfig

class SovythosConfig(PretrainedConfig):
    model_type = "sovythos"

    def __init__(
        self,
        vocab_size=32000,
        dim=512,
        heads=16,
        layers=12,
        block_size=512,
        **kwargs
    ):
        super().__init__(**kwargs)
        self.vocab_size = vocab_size
        self.dim = dim
        self.heads = heads
        self.layers = layers
        self.block_size = block_size