ru-toxicity-encoder

Multi-Task Toxicity Classifier (Russian)

Модель для одновременного определения трёх типов токсичности в русскоязычных текстах:

• profanity — ненормативная лексика
• threat — угрозы
• illegal — призывы к нарушению закона

Архитектура

Модель основана на cointegrated/rubert-tiny2 с тремя независимыми классификационными головами (Multi-Task Learning).

Обучение

Датасет: Gloomreach/ru-toxicity-int-ml

Метрики качества

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

Через 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%}")