Hugging Face's logo

IvanFed
/
rubert-toxicity-multitask

Text Classification
PyTorch
Russian
russian
toxicity
profanity
threat-detection
multi-task
bert
rubert
Model card Files
xet
 Community
rubert-toxicity-multitask
File size: 6,061 Bytes
eb0808f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
---
language: ru
license: apache-2.0
tags:
- russian
- toxicity
- profanity
- threat-detection
- multi-task
- bert
- rubert
- text-classification
datasets:
- IvanFed/russian-toxic-comments-multilabel
---

# RuBERT Multi-Task Toxicity Classifier

## Model Description

Эта модель является **Multi-Task классификатором** токсичности на основе легковесного энкодера `cointegrated/rubert-tiny2`. Она обучена одновременно предсказывать три класса токсичности:

1. **Profanity** (ненормативная лексика) — мат, оскорбления, нецензурная брань
2. **Threat** (угрозы) — явные или скрытые угрозы
3. **Illegal** (запросы на нарушение закона) — запросы о незаконных действиях (в данном датасете отсутствует)

Модель имеет три независимые головы (линейные слоя), каждая из которых выдает логит для своего класса. Вероятности получаются применением сигмоиды.

## Training Data

Модель обучена на датасете [Russian Toxic Comments Multi-label](https://huggingface.co/datasets/IvanFed/russian-toxic-comments-multilabel), содержащем **248,288** размеченных комментариев из социальной сети ok.ru.

## Metrics (on Validation Set)

| Class | Threshold | Precision | Recall | F1-Score |
|-------|-----------|-----------|--------|----------|
| Profanity | 0.45 | 0.8921 | 0.8969 | 0.8945 |
| Threat | 0.35 | 0.7654 | 0.7998 | 0.7823 |
| Illegal | 0.50 | 0.0000 | 0.0000 | 0.0000 |

*Примечание:* Класс `illegal` отсутствует в датасете, поэтому метрики равны нулю.

## How to Use

### 1. Установка зависимостей

```bash
pip install torch transformers huggingface-hub
```

### 2. Загрузка модели и токенизатора

```python
import torch
from transformers import AutoTokenizer
import json

tokenizer = AutoTokenizer.from_pretrained("IvanFed/rubert-toxicity-multitask")

with open("thresholds.json", "r") as f:
    thresholds = json.load(f)

class MultiTaskToxicityEncoder(torch.nn.Module):
    def __init__(self, model_name="cointegrated/rubert-tiny2"):
        super().__init__()
        from transformers import AutoModel
        self.encoder = AutoModel.from_pretrained(model_name)
        self.hidden_size = self.encoder.config.hidden_size
        self.dropout = torch.nn.Dropout(0.1)
        self.profanity_head = torch.nn.Linear(self.hidden_size, 1)
        self.threat_head = torch.nn.Linear(self.hidden_size, 1)
        self.illegal_head = torch.nn.Linear(self.hidden_size, 1)

    def forward(self, input_ids, attention_mask):
        outputs = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
        cls_embedding = outputs.last_hidden_state[:, 0, :]
        cls_embedding = self.dropout(cls_embedding)
        return (
            self.profanity_head(cls_embedding),
            self.threat_head(cls_embedding),
            self.illegal_head(cls_embedding)
        )

model = MultiTaskToxicityEncoder()
model.load_state_dict(torch.load("pytorch_model.bin", map_location="cpu"))
model.eval()

def predict_toxicity(text, threshold_dict=thresholds):
    encoded = tokenizer(text, truncation=True, padding='max_length', max_length=128, return_tensors='pt')
    with torch.no_grad():
        profanity_logits, threat_logits, illegal_logits = model(encoded['input_ids'], encoded['attention_mask'])
        profanity_prob = torch.sigmoid(profanity_logits).item()
        threat_prob = torch.sigmoid(threat_logits).item()
        illegal_prob = torch.sigmoid(illegal_logits).item()
    return {
        'profanity': {'probability': profanity_prob, 'prediction': profanity_prob >= threshold_dict.get('Profanity', 0.5)},
        'threat': {'probability': threat_prob, 'prediction': threat_prob >= threshold_dict.get('Threat', 0.5)},
        'illegal': {'probability': illegal_prob, 'prediction': illegal_prob >= threshold_dict.get('Illegal', 0.5)}
    }

text = "Ты дурак!"
print(predict_toxicity(text))
```

## Training Details

- **Base Model:** `cointegrated/rubert-tiny2`
- **Batch Size:** 64
- **Learning Rate:** 2e-5
- **Optimizer:** AdamW
- **Scheduler:** ReduceLROnPlateau (factor=0.5, patience=2)
- **Early Stopping:** patience=3 (остановка на 3-й эпохе)
- **Max Sequence Length:** 128 токенов
- **Loss Function:** BCEWithLogitsLoss (сумма потерь для трех голов)
- **Number of Epochs Trained:** 6 (best model at epoch 3)

## Limitations

- Класс `illegal` отсутствует в обучающем датасете, поэтому модель не может его обнаружить.
- Модель обучена только на русскоязычных комментариях из социальной сети ok.ru, может плохо работать на других доменах.
- Для улучшения качества на классе `threat` рекомендуется собрать дополнительные данные.

## Acknowledgements

Благодарность авторам оригинального датасета [Toxic Russian Comments](https://www.kaggle.com/datasets/alexandersemiletov/toxic-russian-comments) за предоставленные данные.

## Citation

Если вы используете эту модель в своей работе, пожалуйста, ссылайтесь на:

```
@misc{rubert_toxicity_multitask,
  author = {IvanFed},
  title = {RuBERT Multi-Task Toxicity Classifier},
  year = {2026},
  publisher = {Hugging Face},
  howpublished = {\url{https://huggingface.co/IvanFed/rubert-toxicity-multitask}}
}
```