| 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)} | |