ru-multitask-toxicity-encoder

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

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

Пороги подобраны индивидуально для каждого класса по максимуму 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. Создано в учебных целях для курса по трансформерным энкодерам.