File size: 4,594 Bytes
f2cd704
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
7118a79
2467616
 
f2cd704
 
 
 
 
 
 
 
 
1c11c87
f2cd704
1c11c87
f2cd704
 
f32a21a
f2cd704
 
 
3e87598
 
 
 
 
 
f2cd704
 
f32a21a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
f2cd704
 
 
 
 
 
 
 
2467616
f2cd704
 
2467616
f2cd704
 
2467616
f2cd704
 
 
 
 
 
 
 
2467616
 
 
 
 
 
 
 
 
 
 
 
 
 
 
f2cd704
 
 
 
 
 
 
 
 
f32a21a
 
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
141
142
143
144
145
146
147
148
149
150

---
language: ru
license: mit
tags:
- toxicity
- multilabel
- russian
- multitask
- rubert
- profanity
- threat
- illegal
metrics:
- f1
- precision
- recall
- auc
pipeline_tag: text-classification
---

# Multi-Task Toxicity Classifier

## Описание

Модель для одновременного обнаружения трёх типов токсичности в русскоязычных текстах:
- **Profanity** (ненормативная лексика)
- **Threat** (угрозы)
- **Illegal** (запросы на нарушение закона)

Модель основана на энкодере `cointegrated/rubert-tiny2` с тремя независимыми классификационными головами.

## Метрики на валидационной выборке

| Класс | Threshold | Precision | Recall | F1-Score |
|-------|-----------|-----------|--------|----------|
| Profanity | 0.51 | 0.9048 | 0.9429 | 0.9235 |
| Threat | 0.35 | 0.7462 | 0.8097 | 0.7766 |
| Illegal | 0.25 | 0.6038 | 0.6598 | 0.6305 |

**Macro F1-Score:** 0.7769


## Использование

```python
from transformers import AutoTokenizer, AutoModel
import torch
import torch.nn as nn
import json
import requests

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

tokenizer = AutoTokenizer.from_pretrained("qquarkq/multitask-toxicity-classifier")

# Загрузка конфигурации с порогами
url = "https://huggingface.co/qquarkq/multitask-toxicity-classifier/resolve/main/config.json"
response = requests.get(url)
config = response.json()
#Или
#with open("config.json", "r") as f:
#    config = json.load(f)
thresholds = config["thresholds"]

class MultiTaskToxicityEncoder(nn.Module):
    def __init__(self, model_name, dropout=0.2, freeze_encoder=False):
        super().__init__()

        self.encoder = AutoModel.from_pretrained(model_name)
        if freeze_encoder:
          for param in self.encoder.parameters():
              param.requires_grad = False
        self.hidden_size = self.encoder.config.hidden_size
        self.dropout = nn.Dropout(dropout)

        #Головы
        self.profanity_head = nn.Linear(self.hidden_size, 1)  # Нецензурная лексика
        self.threat_head = nn.Linear(self.hidden_size, 1)     # Угрозы
        self.illegal_head = nn.Linear(self.hidden_size, 1) # Незаконный контент



    def forward(self, input_ids, attention_mask):
        outputs = self.encoder(
            input_ids=input_ids,
            attention_mask=attention_mask,
            return_dict=True
        )
        cls_embedding = outputs.last_hidden_state[:, 0, :]
        cls_embedding = self.dropout(cls_embedding)

        profanity_logits = self.profanity_head(cls_embedding)
        threat_logits = self.threat_head(cls_embedding)
        illegal_logits = self.illegal_head(cls_embedding)

        return profanity_logits, threat_logits, illegal_logits

model = MultiTaskToxicityEncoder(model_name=config["model_name"])

def predict_toxicity(text):
    encoding = tokenizer(
        text,
        truncation=True,
        padding='max_length',
        max_length=128,
        return_tensors='pt'
    )

    with torch.no_grad():
        profanity_logits, threat_logits, illegal_logits = model(
            encoding['input_ids'],
            encoding['attention_mask']
        )

    profanity_prob = torch.sigmoid(profanity_logits).item()
    threat_prob = torch.sigmoid(threat_logits).item()
    illegal_prob = torch.sigmoid(illegal_logits).item()

    profanity_pred = int(profanity_prob >= thresholds['profanity'])
    threat_pred = int(threat_prob >= thresholds['threat'])
    illegal_pred = int(illegal_prob >= thresholds['illegal'])
    return {
        'profanity': {
            'probability': profanity_prob,
            'prediction': profanity_pred,
            'confidence': profanity_prob * 100
        },
        'threat': {
            'probability': threat_prob,
            'prediction': threat_pred,
            'confidence': threat_prob * 100
        },
        'illegal': {
            'probability': illegal_prob,
            'prediction': illegal_pred,
            'confidence': illegal_prob * 100
        }
    }

# Пример использования
text = "Ты полный идиот!"
result = predict_toxicity(text)
print(result)
````
Датасет
Модель обучена на датасете [qquarkq/russian-toxic-multilabel-comments](https://huggingface.co/datasets/qquarkq/russian-toxic-multilabel-comments)