rubert-toxicity-multitask

RuBERT Multi-Task Toxicity Classifier

Model Description

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

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

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

Training Data

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

Metrics (on Validation Set)

Примечание: Класс 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}}
}