README.md

---
language:
- ru
license: mit
tags:
- text-classification
- toxicity
- multi-label-classification
- multi-task-learning
- russian
- bert
base_model: cointegrated/rubert-tiny2
---

# RU Multi-Task Toxicity Encoder

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

## Описание

Модель построена на базе [cointegrated/rubert-tiny2](https://huggingface.co/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` при обучении.

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

```python
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](https://huggingface.co/datasets/AtesiT/ru-multitask-toxicity)
- **Функция потерь**: `BCEWithLogitsLoss` с `pos_weight` для каждого класса
- **Оптимизатор**: AdamW (lr=2e-5, weight_decay=0.01)
- **Early Stopping**: patience=3 по валидационному лоссу
- **Макс. длина последовательности**: 64 токена

## Ограничения

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

## Лицензия

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