fdlvaaa/rubert-tiny2-multitask-toxicity

Multi-Task Toxicity Classification Model (ruBERT-tiny2)

Данная модель предназначена для автоматической модерации русскоязычных текстов и одновременно решает 3 задачи классификации:

1. Profanity — ненормативная лексика и мат.
2. Threat — угрозы насилия и физической расправы.
3. Illegal — обсуждение/призывы к противоправным и запрещенным действиям.

Основана на лёгком энкодере cointegrated/rubert-tiny2 с тремя независимыми классификационными головами.

📊 Метрики на Валидации

Категория	Precision	Recall	F1-Score	Оптимальный порог
Profanity	0.8484	0.8534	0.8509	0.90
Threat	0.6532	0.8526	0.7397	0.85
Illegal	0.0000	0.0000	0.0000	0.50

Итоговый Macro F1-Score: 0.5302

🚀 Пример использования (Python / PyTorch)

import torch
import torch.nn as nn
from huggingface_hub import hf_hub_download
from transformers import AutoModel, AutoTokenizer


class MultiTaskToxicityEncoder(nn.Module):

    def __init__(self, model_name="cointegrated/rubert-tiny2"):
        super().__init__()
        self.encoder = AutoModel.from_pretrained(model_name)
        hidden_size = self.encoder.config.hidden_size
        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):
        outputs = self.encoder(
            input_ids=input_ids, attention_mask=attention_mask
        )
        cls_embedding = outputs.last_hidden_state[:, 0, :]
        return {
            "profanity": self.profanity_head(cls_embedding).squeeze(-1),
            "threat": self.threat_head(cls_embedding).squeeze(-1),
            "illegal": self.illegal_head(cls_embedding).squeeze(-1),
        }


# Загрузка токенизатора и весов
repo_id = "fdlvaaa/rubert-tiny2-multitask-toxicity"
tokenizer = AutoTokenizer.from_pretrained(repo_id)
model = MultiTaskToxicityEncoder()

weights_path = hf_hub_download(repo_id=repo_id, filename="model.pt")
model.load_state_dict(torch.load(weights_path, map_location="cpu"))
model.eval()

# Инференс
text = "Ты еще пожалеешь, я найду тебя!"
inputs = tokenizer(text, return_tensors="pt")

with torch.no_grad():
    outputs = model(inputs["input_ids"], inputs["attention_mask"])
    probs = {
        k: torch.sigmoid(v).item() for k, v in outputs.items()
    }

print("Вероятности:", probs)