RuBERT Multi-Task Toxicity Classifier

Model Description

Эта модель является Multi-Task классификатором токсичности на основе легковесного энкодера cointegrated/rubert-tiny2. Она обучена одновременно предсказывать три класса токсичности:

Profanity (ненормативная лексика) — мат, оскорбления, нецензурная брань
Threat (угрозы) — явные или скрытые угрозы
Illegal (запросы на нарушение закона) — запросы о незаконных действиях (в данном датасете отсутствует)

Модель имеет три независимые головы (линейные слоя), каждая из которых выдает логит для своего класса. Вероятности получаются применением сигмоиды.

Training Data

Модель обучена на датасете Russian Toxic Comments Multi-label, содержащем 248,288 размеченных комментариев из социальной сети ok.ru.

Metrics (on Validation Set)

Class	Threshold	Precision	Recall	F1-Score
Profanity	0.45	0.8921	0.8969	0.8945
Threat	0.35	0.7654	0.7998	0.7823
Illegal	0.50	0.0000	0.0000	0.0000

Примечание: Класс illegal отсутствует в датасете, поэтому метрики равны нулю.

How to Use

1. Установка зависимостей

pip install torch transformers huggingface-hub

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

import torch
from transformers import AutoTokenizer
import json

tokenizer = AutoTokenizer.from_pretrained("IvanFed/rubert-toxicity-multitask")

with open("thresholds.json", "r") as f:
    thresholds = json.load(f)

class MultiTaskToxicityEncoder(torch.nn.Module):
    def __init__(self, model_name="cointegrated/rubert-tiny2"):
        super().__init__()
        from transformers import AutoModel
        self.encoder = AutoModel.from_pretrained(model_name)
        self.hidden_size = self.encoder.config.hidden_size
        self.dropout = torch.nn.Dropout(0.1)
        self.profanity_head = torch.nn.Linear(self.hidden_size, 1)
        self.threat_head = torch.nn.Linear(self.hidden_size, 1)
        self.illegal_head = torch.nn.Linear(self.hidden_size, 1)

    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.profanity_head(cls_embedding),
            self.threat_head(cls_embedding),
            self.illegal_head(cls_embedding)
        )

model = MultiTaskToxicityEncoder()
model.load_state_dict(torch.load("pytorch_model.bin", map_location="cpu"))
model.eval()

def predict_toxicity(text, threshold_dict=thresholds):
    encoded = tokenizer(text, truncation=True, padding='max_length', max_length=128, return_tensors='pt')
    with torch.no_grad():
        profanity_logits, threat_logits, illegal_logits = model(encoded['input_ids'], encoded['attention_mask'])
        profanity_prob = torch.sigmoid(profanity_logits).item()
        threat_prob = torch.sigmoid(threat_logits).item()
        illegal_prob = torch.sigmoid(illegal_logits).item()
    return {
        'profanity': {'probability': profanity_prob, 'prediction': profanity_prob >= threshold_dict.get('Profanity', 0.5)},
        'threat': {'probability': threat_prob, 'prediction': threat_prob >= threshold_dict.get('Threat', 0.5)},
        'illegal': {'probability': illegal_prob, 'prediction': illegal_prob >= threshold_dict.get('Illegal', 0.5)}
    }

text = "Ты дурак!"
print(predict_toxicity(text))

Training Details

Base Model: cointegrated/rubert-tiny2
Batch Size: 64
Learning Rate: 2e-5
Optimizer: AdamW
Scheduler: ReduceLROnPlateau (factor=0.5, patience=2)
Early Stopping: patience=3 (остановка на 3-й эпохе)
Max Sequence Length: 128 токенов
Loss Function: BCEWithLogitsLoss (сумма потерь для трех голов)
Number of Epochs Trained: 6 (best model at epoch 3)

Limitations

Класс illegal отсутствует в обучающем датасете, поэтому модель не может его обнаружить.
Модель обучена только на русскоязычных комментариях из социальной сети ok.ru, может плохо работать на других доменах.
Для улучшения качества на классе threat рекомендуется собрать дополнительные данные.

Acknowledgements

Благодарность авторам оригинального датасета Toxic Russian Comments за предоставленные данные.

Citation

Если вы используете эту модель в своей работе, пожалуйста, ссылайтесь на:

@misc{rubert_toxicity_multitask,
  author = {IvanFed},
  title = {RuBERT Multi-Task Toxicity Classifier},
  year = {2026},
  publisher = {Hugging Face},
  howpublished = {\url{https://huggingface.co/IvanFed/rubert-toxicity-multitask}}
}

Downloads last month: 28

IvanFed
/

rubert-toxicity-multitask