from transformers import PretrainedConfig

class KiyEngineConfig(PretrainedConfig):
    model_type = "kiyengine"
    
    def __init__(
        self, 
        d_model=384, 
        n_layers=4, 
        n_experts=8, 
        top_k=2, 
        d_state=16, 
        d_conv=4, 
        expansion_factor=2, 
        vocab_size=768, 
        **kwargs
    ):
        self.d_model = d_model
        self.n_layers = n_layers
        self.n_experts = n_experts
        self.top_k = top_k
        self.d_state = d_state
        self.d_conv = d_conv
        self.expansion_factor = expansion_factor
        self.vocab_size = vocab_size
        super().__init__(**kwargs)