RU Multi-Task Toxicity Encoder

Лёгкая multi-task модель для классификации токсичности русскоязычных текстов по трём независимым категориям одновременно.

Описание

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

Каждая голова предсказывает один бинарный класс:

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

Архитектура

Encoder (rubert-tiny2)
    └── [CLS] embedding (312-dim)
            ├── Dropout (p=0.2)
            ├── Linear(312 → 1) → profanity_logit
            ├── Linear(312 → 1) → threat_logit
            └── Linear(312 → 1) → illegal_logit

Метрики (валидационная выборка)

Класс	Порог	Precision	Recall	F1
profanity	0.60	0.9454	0.9256	0.9354
threat	0.75	1.0000	0.9286	0.9630
illegal	0.16	1.0000	1.0000	1.0000

Пороги подобраны индивидуально для каждого класса по максимуму F1-score на валидационной выборке. Классы сильно несбалансированы — это учтено через pos_weight в BCEWithLogitsLoss при обучении.

Пример использования

import json
import torch
import torch.nn as nn
from transformers import AutoModel, AutoTokenizer

# 1. Загрузка токенизатора и конфига
tokenizer = AutoTokenizer.from_pretrained("AtesiT/ru-multitask-toxicity-encoder")

with open("toxicity_config.json") as f:
    config = json.load(f)

# 2. Определение архитектуры
class MultiTaskToxicityEncoder(nn.Module):
    def __init__(self, model_name, hidden_size, dropout=0.2):
        super().__init__()
        self.encoder = AutoModel.from_pretrained(model_name)
        self.dropout = nn.Dropout(dropout)
        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):
        out = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
        cls = self.dropout(out.last_hidden_state[:, 0, :])
        return (
            self.profanity_head(cls),
            self.threat_head(cls),
            self.illegal_head(cls),
        )

# 3. Загрузка весов
model = MultiTaskToxicityEncoder(
    model_name=config["base_model"],
    hidden_size=config["hidden_size"],
)
state_dict = torch.load("model_weights.pt", map_location="cpu")
model.load_state_dict(state_dict)
model.eval()

# 4. Инференс
thresholds = config["thresholds"]

def predict(text):
    enc = tokenizer(
        text, return_tensors="pt",
        padding="max_length", truncation=True,
        max_length=config["max_length"],
    )
    with torch.no_grad():
        p_logit, t_logit, i_logit = model(
            enc["input_ids"], enc["attention_mask"]
        )
    probs = {
        "profanity": torch.sigmoid(p_logit).item(),
        "threat":    torch.sigmoid(t_logit).item(),
        "illegal":   torch.sigmoid(i_logit).item(),
    }
    labels = {k: int(v >= thresholds[k]) for k, v in probs.items()}
    return {"probs": probs, "labels": labels}

print(predict("Ты полный идиот, заткнись!"))

Обучение

Датасет: AtesiT/ru-multitask-toxicity
Функция потерь: BCEWithLogitsLoss с pos_weight для каждого класса
Оптимизатор: AdamW (lr=2e-5, weight_decay=0.01)
Early Stopping: patience=3 по валидационному лоссу
Макс. длина последовательности: 64 токена

Ограничения

Модель обучена на синтетически аугментированных данных для редких классов (threat, illegal) — качество на реальных данных может отличаться.
Короткие тексты (< 5 слов) могут обрабатываться менее точно.
Модель не учитывает контекст диалога — только отдельное сообщение.

Лицензия

MIT. Создано в учебных целях для курса по трансформерным энкодерам.

AtesiT/ru-multitask-toxicity-encoder