from transformers import PretrainedConfig

class ModernDenseConfig(PretrainedConfig):
    model_type='modern_dense_mha_gated_ffn_router'
    def __init__(self,dense_config=None,**kwargs):
        super().__init__(**kwargs)
        self.dense_config=dense_config or {}
        self.vocab_size=int(self.dense_config.get('vocab',16384))
        self.hidden_size=int(self.dense_config.get('width',256))
        self.tie_word_embeddings=False
        self.auto_map={
            'AutoConfig':'configuration_dense.ModernDenseConfig',
            'AutoModel':'modeling_dense.ModernDenseModel',
            'AutoModelForCausalLM':'modeling_dense.ModernDenseForCausalLM',
        }