pogram1st/multitask-toxicity-comments

Multi-Task Toxicity RU

Модель для классификации токсичных комментариев на русском языке. Модель использует архитектуру RuBERT-tiny2 с тремя независимыми «головами» классификации для одновременного выявления разных типов нарушений.

Метрики производительности

Модель была обучена и оптимизирована по порогам для каждого класса:

Класс	Оптимальный порог	Precision	Recall	F1-score
Ненормативная лексика	0.75	0.9402	0.9495	0.9448
Угрозы	0.80	0.8031	0.8421	0.8221
Запросы на нарушение закона	0.90	0.9242	0.9531	0.9385

Быстрый старт

import torch
from transformers import AutoTokenizer, AutoModel

REPO_ID = "pogram1st/multitask-toxicity-comments"  
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

tokenizer = AutoTokenizer.from_pretrained(REPO_ID)
model = AutoModel.from_pretrained(REPO_ID, trust_remote_code=True)

model.to(DEVICE)
model.eval()

def predict_toxicity(text):
    encoding = tokenizer(
        text,
        max_length=64,
        padding='max_length',
        truncation=True,
        return_tensors='pt'
    ).to(DEVICE)
    
    with torch.no_grad():
        out_prof, out_threat, out_ill = model(encoding['input_ids'], encoding['attention_mask'])
        
    prob_prof = torch.sigmoid(out_prof).item()
    prob_threat = torch.sigmoid(out_threat).item()
    prob_ill = torch.sigmoid(out_ill).item()
    
    print(f"Текст: '{text}'")
    print(f"Ненормативная лексика: {prob_prof:.0%} -> {'Запрещенный контент' if prob_prof >= 0.75 else 'Ок'}")
    print(f"Угроза: {prob_threat:.0%} -> {'Запрещенный контент' if prob_threat >= 0.80 else 'Ок'}")
    print(f"Запрос на нарушение закона: {prob_ill:.0%} -> {'Запрещенный контент' if prob_ill >= 0.90 else 'Ок'}\n")

predict_toxicity("Это приятно знать, но тем не менее.")