from transformers import PretrainedConfig

class NeuroCLRConfig(PretrainedConfig):
    model_type = "neuroclr"

    def __init__(
        self,
        TSlength: int = 128,
        nhead: int = 2,
        nlayer: int = 2,
        projector_out1: int = 128,
        projector_out2: int = 64,

        # classification
        num_labels: int = 2,

        # pooling to avoid flatten dimension mismatch
        pooling: str = "flatten",  # "mean" recommended; "flatten" only if seq_len==1

        normalize_input: bool = True,
        **kwargs
    ):
        super().__init__(**kwargs)
        self.TSlength = TSlength
        self.nhead = nhead
        self.nlayer = nlayer
        self.projector_out1 = projector_out1
        self.projector_out2 = projector_out2
        self.num_labels = num_labels
        self.pooling = pooling
        self.normalize_input = normalize_input