"""Configuration class for JetonCount."""

from transformers import PretrainedConfig


class JetonCountConfig(PretrainedConfig):
    model_type = "jetoncount"

    def __init__(
        self,
        feature_names=None,
        base_feature_names=None,
        base_feature_dim=7,
        feature_dim=19,
        num_layers=8,
        hidden_dim=32,
        dropout=0.005,
        activation="silu",
        standardize_features=True,
        use_engineered_features=True,
        use_log1p_features=True,
        use_log1p_target=False,
        center_target=False,
        target_offset=0.0,
        feature_mean=None,
        feature_std=None,
        **kwargs,
    ):
        self.base_feature_names = base_feature_names or [
            "chars",
            "words",
            "avg_chars_per_word",
            "punctuation_ratio",
            "symbol_ratio",
            "longest_word_chars",
            "vocab_size",
        ]
        self.feature_names = feature_names or list(self.base_feature_names)
        self.base_feature_dim = int(base_feature_dim)
        self.feature_dim = int(feature_dim)
        self.num_layers = int(num_layers)
        self.hidden_dim = int(hidden_dim)
        self.dropout = float(dropout)
        self.activation = str(activation)
        self.standardize_features = bool(standardize_features)
        self.use_engineered_features = bool(use_engineered_features)
        self.use_log1p_features = bool(use_log1p_features)
        self.use_log1p_target = bool(use_log1p_target)
        self.center_target = bool(center_target)
        self.target_offset = float(target_offset)
        self.feature_mean = feature_mean
        self.feature_std = feature_std
        super().__init__(**kwargs)