"""Configuration for PC-MLP. Required for auto_map + trust_remote_code."""
from transformers import PretrainedConfig

class PCMLPConfig(PretrainedConfig):
    model_type = "pcmlp"

    def __init__(
        self,
        vocab_size=16,
        seq_len=16,
        num_classes=2,
        hidden_dim=32,
        local_loss_weight=0.1,
        **kwargs,
    ):
        self.vocab_size = vocab_size
        self.seq_len = seq_len
        self.num_classes = num_classes
        self.hidden_dim = hidden_dim
        self.local_loss_weight = local_loss_weight
        super().__init__(**kwargs)