from transformers import PretrainedConfig


class ActEstimatorConfig(PretrainedConfig):
    model_type = "ACT-Estimator"

    def __init__(
        self,
        input_shape=(3, 44, 224, 224),
        num_classes=9,
        max_seq_len=44,
        timestamp_dim=1,
        d_model=512,
        num_heads=8,
        dropout=0.1,
        feature_map_size=4,
        **kwargs
    ):
        self.input_shape = input_shape
        self.num_classes = num_classes
        self.max_seq_len = max_seq_len
        self.timestamp_dim = timestamp_dim
        self.d_model = d_model
        self.num_heads = num_heads
        self.dropout = dropout
        self.feature_map_size = feature_map_size
        super().__init__(**kwargs)