from transformers import PretrainedConfig

class Evo2ExonConfig(PretrainedConfig):
    model_type = "evo2_exon_classifier"

    def __init__(self,
                 embedding_dim: int = 8192,   # match your input width
                 hidden_dim:    int = 1024,   # width of hidden layers
                 num_hidden_layers: int = 1,  # depth ≥1
                 **kwargs):
        super().__init__(**kwargs)
        self.embedding_dim      = embedding_dim
        self.hidden_dim         = hidden_dim
        self.num_hidden_layers  = num_hidden_layers