Multi-Task Toxicity Classifier (Russian)
Модель для одновременного определения трёх типов токсичности в русскоязычных текстах:
- profanity — ненормативная лексика
- threat — угрозы
- illegal — призывы к нарушению закона
Архитектура
Модель основана на cointegrated/rubert-tiny2 с тремя независимыми классификационными головами (Multi-Task Learning).
Обучение
Датасет: Gloomreach/ru-toxicity-int-ml
Метрики качества
| Класс | Precision | Recall | F1-Score | Threshold |
|---|---|---|---|---|
| profanity | 0.939856 | 0.961855 | 0.950728 | 0.2 |
| threat | 0.956522 | 0.956522 | 0.956522 | 0.4 |
| illegal | 1.000000 | 0.600000 | 0.750000 | 0.5 |
Использование
Через transformers
import torch
import torch.nn as nn
from transformers import AutoModel, AutoTokenizer
class MultiTaskToxicityEncoder(nn.Module):
def __init__(self, model_name="cointegrated/rubert-tiny2"):
super().__init__()
self.encoder = AutoModel.from_pretrained(model_name)
hidden_size = self.encoder.config.hidden_size
self.profanity_head = nn.Linear(hidden_size, 1)
self.threat_head = nn.Linear(hidden_size, 1)
self.illegal_head = nn.Linear(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, :]
return (
self.profanity_head(cls_embedding),
self.threat_head(cls_embedding),
self.illegal_head(cls_embedding)
)
#загрузка модели
model = MultiTaskToxicityEncoder.from_pretrained("Gloomreach/ru-toxicity-multi-task-encoder")
tokenizer = AutoTokenizer.from_pretrained("cointegrated/rubert-tiny2")
text = "Пример текста для проверки"
inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=128)
with torch.no_grad():
logits = model(inputs["input_ids"], inputs["attention_mask"])
probs = torch.sigmoid(torch.cat(logits, dim=1)).numpy()[0]
print(f"Profanity: {probs[0]:.2%}")
print(f"Threat: {probs[1]:.2%}")
print(f"Illegal: {probs[2]:.2%}")