from transformers import PretrainedConfig, AutoConfig

class MLPConfig(PretrainedConfig):
    r"""
    Config for the MLP model.

    Args:
        embedding_size (int, 8): Size of the input embeddings (last dimension of the 3D input).
        sequence_length (int, 8): Number of tokens in input sequence (middle dimension; must be fixed).
        num_labels (int, 32): Number of output labels (for multi-label classification).
        hidden_size (int, 2048): Size of each hidden layer.
        num_hidden_layers (int, 3): How many hidden layers to stack.
        dropout (float, 0.1): Dropout probability for hidden layers.
    """
    
    model_type = "mlp"

    def __init__(
        self,
        embedding_size=8,
        sequence_length=8,
        num_labels=32,
        hidden_size=2048,
        num_hidden_layers=3,
        dropout=0.1,
        **kwargs
    ):
        super().__init__(**kwargs)
        self.embedding_size = embedding_size
        self.sequence_length = sequence_length
        self.num_labels = num_labels
        self.hidden_size = hidden_size
        self.num_hidden_layers = num_hidden_layers
        self.dropout = dropout


AutoConfig.register("mlp", MLPConfig)