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

Intended Use

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

• фильтрации нецензурной брани;
• выделения угроз для приоритетной проверки модераторами;
• обнаружения сообщений, содержащих запросы на совершение незаконных действий.

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

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

Training Data

Модель обучена на датасете, собранном из двух открытых источников:

• petaevd/russian-toxicity-multilabel – для класса illegal
• AlexSham/Toxic_Russian_Comments – для классов profanity и threat

Датасет был очищен, сбалансирован и опубликован отдельно:
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

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

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)