README.md

---
language: ru
license: mit
tags:
- toxicity
- multilabel-classification
- russian
- rubert-tiny2
- multitask-learning
pretty_name: RuBERT Tiny2 Toxicity Multitask
model_name: rubert-tiny2-toxicity-multitask
---

# Model Card for rubert-tiny2-toxicity-multitask

## Model Details

**Model Description:**  
Эта модель представляет собой многозадачный классификатор токсичности для русского языка. Она построена на основе лёгкого энкодера `cointegrated/rubert-tiny2` и содержит три независимые классификационные «головы» для одновременного предсказания трёх классов:

- **profanity** – нецензурная лексика и оскорбления
- **threat** – угрозы
- **illegal** – запросы, связанные с незаконными действиями

Модель обучена в мультитаск-режиме (Multi-Task Learning) с использованием взвешенной функции потерь для учёта дисбаланса классов.

**Разработчик:** Arrtemwolf  
**Модель доступна на Hugging Face:** [https://huggingface.co/Arrtemwolf/rubert-tiny2-toxicity-multitask](https://huggingface.co/Arrtemwolf/rubert-tiny2-toxicity-multitask)

## Intended Use

**Назначение:**  
Модель предназначена для автоматической модерации текстовых сообщений в социальных сетях, чатах, комментариях. Она может использоваться для:
- фильтрации нецензурной брани;
- выделения угроз для приоритетной проверки модераторами;
- обнаружения сообщений, содержащих запросы на совершение незаконных действий.

**Ограничения:**  
- Модель обучена на ограниченном наборе данных (около 5500 примеров) и может не охватывать все разнообразие русскоязычной токсичности.
- Рекомендуется использовать в связке с дополнительными фильтрами и человеческой модерацией для критичных применений.

## Training Data

Модель обучена на датасете, собранном из двух открытых источников:
- [petaevd/russian-toxicity-multilabel](https://huggingface.co/datasets/petaevd/russian-toxicity-multilabel) – для класса `illegal`
- [AlexSham/Toxic_Russian_Comments](https://huggingface.co/datasets/AlexSham/Toxic_Russian_Comments) – для классов `profanity` и `threat`

Датасет был очищен, сбалансирован и опубликован отдельно:  
[Arrtemwolf/ru-toxic-multi-label-clean](https://huggingface.co/datasets/Arrtemwolf/ru-toxic-multi-label-clean)

## Evaluation Results

Модель оценивалась на валидационной выборке (20% от общего датасета). Для каждого класса был подобран оптимальный порог вероятности, максимизирующий **F1-меру**.

| Класс      | Порог | Precision | Recall | F1-score |
|------------|-------|-----------|--------|----------|
| profanity  | 0.50  | 0.821     | 0.908  | **0.862**|
| threat     | 0.15  | 0.261     | 0.571  | **0.358**|
| illegal    | 0.40  | 1.000     | 1.000  | **1.000**|

> Примечание: высокий F1 для класса `illegal` обусловлен небольшим количеством положительных примеров в валидации; на более крупных выборках результаты могут отличаться.

## How to Use

### Загрузка модели и токенизатора

```python
from transformers import AutoTokenizer, AutoModel
import torch

MODEL_NAME = "Arrtemwolf/rubert-tiny2-toxicity-multitask"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
encoder = AutoModel.from_pretrained(MODEL_NAME)

# Загрузка обученных голов (классификаторов)
# Внимание: головы сохранены отдельно, их нужно загрузить и прикрепить к модели
# Ниже приведён пример класса-обёртки, который можно использовать после загрузки весов голов.
Для удобства рекомендуется использовать класс MultiTaskToxicityEncoder, который объединяет энкодер и три головы. Веса голов сохранены в файле multitask_heads.pt в репозитории. Пример загрузки:

python
class MultiTaskToxicityEncoder(torch.nn.Module):
    def __init__(self, encoder):
        super().__init__()
        self.encoder = encoder
        hidden_size = encoder.config.hidden_size
        self.head_profanity = torch.nn.Linear(hidden_size, 1)
        self.head_threat = torch.nn.Linear(hidden_size, 1)
        self.head_illegal = torch.nn.Linear(hidden_size, 1)
        self.dropout = torch.nn.Dropout(0.3)

    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, :]
        cls_embedding = self.dropout(cls_embedding)
        return (self.head_profanity(cls_embedding),
                self.head_threat(cls_embedding),
                self.head_illegal(cls_embedding))

# Загружаем энкодер
encoder = AutoModel.from_pretrained(MODEL_NAME)
model = MultiTaskToxicityEncoder(encoder)

# Загружаем веса голов
state_dict = torch.load("multitask_heads.pt", map_location="cpu")
model.load_state_dict(state_dict, strict=False)  # strict=False, т.к. веса только для голов

model.eval()
Предсказание для одного текста
python
def predict(text, model, tokenizer, device="cpu"):
    encoded = tokenizer(text, padding=True, truncation=True, max_length=256, return_tensors="pt")
    input_ids = encoded["input_ids"].to(device)
    attention_mask = encoded["attention_mask"].to(device)

    with torch.no_grad():
        logit_p, logit_t, logit_i = model(input_ids, attention_mask)
        prob_p = torch.sigmoid(logit_p).item()
        prob_t = torch.sigmoid(logit_t).item()
        prob_i = torch.sigmoid(logit_i).item()

    # Пороги (оптимальные, полученные на валидации)
    thresholds = {"profanity": 0.50, "threat": 0.15, "illegal": 0.40}
    return {
        "profanity": {"prob": prob_p, "label": prob_p >= thresholds["profanity"]},
        "threat":    {"prob": prob_t, "label": prob_t >= thresholds["threat"]},
        "illegal":   {"prob": prob_i, "label": prob_i >= thresholds["illegal"]},
    }

# Пример
text = "Ты мне угрожаешь? Я вызову полицию!"
result = predict(text, model, tokenizer)
print(result)