from transformers import PretrainedConfig


class CustomConvNextConfig(PretrainedConfig):
    model_type = "custom_convnext"

    def __init__(
        self,
        input_channels: int = 3,
        num_classes: int = 3,
        device: str = "cpu",
        convnext_mult: int = 2,
        layer_norm: bool = False,  # Fixed Bool to bool
        dropout: float = 0.0,
        pooling_choice: str = "MaxPool2d",
        activation_choice: str = "ReLU",
        **kwargs  # Add a comma before **kwargs
    ):

        # Initialize class attributes
        self.input_channels = input_channels
        self.num_classes = num_classes
        self.device = device
        self.convnext_mult = convnext_mult
        self.layer_norm = layer_norm
        self.dropout = dropout
        self.pooling_choice = pooling_choice
        self.activation_choice = activation_choice

        # Call the super class's constructor
        super().__init__(**kwargs)