from transformers import PreTrainedConfig

class MultiTaskToxicityConfig(PreTrainedConfig):
    model_type = "multitask_toxicity"
    def __init__(self, encoder_name="cointegrated/rubert-tiny2", labels=None,
                 thresholds=None, max_length=128, dropout=0.15, **kwargs):
        super().__init__(**kwargs)
        self.encoder_name = encoder_name
        self.labels = labels or ["profanity", "threat", "illegal"]
        self.thresholds = thresholds or {label: 0.5 for label in self.labels}
        self.max_length = int(max_length)
        self.dropout = float(dropout)
        self.num_labels = len(self.labels)
        self.id2label = {i: label for i, label in enumerate(self.labels)}
        self.label2id = {label: i for i, label in enumerate(self.labels)}