configuration_mlp.py

"""Configuration for a simple tanh MLP.

This file is self-contained: it is uploaded to the Hub and executed at load
time via `trust_remote_code=True`. It must not import any project-local code.
"""

try:  # transformers >= 5
    from transformers import PreTrainedConfig
except ImportError:  # transformers < 5
    from transformers import PretrainedConfig as PreTrainedConfig


class MLPConfig(PreTrainedConfig):
    # `model_type` is what wires the config into the AutoConfig registry.
    model_type = "tanh_mlp"

    def __init__(
        self,
        input_dim: int = 16,
        hidden_dim: int = 32,
        output_dim: int = 4,
        num_hidden_layers: int = 2,
        **kwargs,
    ):
        self.input_dim = input_dim
        self.hidden_dim = hidden_dim
        self.output_dim = output_dim
        self.num_hidden_layers = num_hidden_layers
        # PreTrainedConfig defines many extra fields; accept and forward kwargs.
        super().__init__(**kwargs)