from transformers import PretrainedConfig
from typing import List

class OffensivenessEstimationConfig(PretrainedConfig):
    model_type = "offensiveness_estimation"

    def __init__(
        self,
        language_model: str   = 'Twitter/twhin-bert-large',
        output_class_num: int = 11,
        reinit_n_layers: int  = 1,
        dropout_rate: float   = 0.1,
        **kwargs,
    ):
        # if block_type not in ["basic", "bottleneck"]:
        #     raise ValueError(f"`block_type` must be 'basic' or bottleneck', got {block_type}.")
        # if stem_type not in ["", "deep", "deep-tiered"]:
        #     raise ValueError(f"`stem_type` must be '', 'deep' or 'deep-tiered', got {stem_type}.")

        self.language_model   = language_model
        self.output_class_num = output_class_num
        self.reinit_n_layers  = reinit_n_layers
        self.dropout_rate     = dropout_rate
        super().__init__(**kwargs)