---
language: ru
license: mit
tags:
- toxicity
- multilabel
- russian
- multitask
- rubert
- profanity
- threat
- illegal
metrics:
- f1
- precision
- recall
- auc
pipeline_tag: text-classification
---

# Multi-Task Toxicity Classifier

## Описание

Модель для одновременного обнаружения трёх типов токсичности в русскоязычных текстах:
- **Profanity** (ненормативная лексика)
- **Threat** (угрозы)
- **Illegal** (запросы на нарушение закона)

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

## Метрики на валидационной выборке

| Класс | Threshold | Precision | Recall | F1-Score |
|-------|-----------|-----------|--------|----------|
| Profanity | 0.63 | 0.9048 | 0.9429 | 0.9235 |
| Threat | 0.74 | 0.7462 | 0.8097 | 0.7766 |
| Illegal | 0.89 | 0.6038 | 0.6598 | 0.6305 |

**Macro F1-Score:** 0.7769


## Использование

```python
from transformers import AutoTokenizer, AutoModel
import torch
import json

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

tokenizer = AutoTokenizer.from_pretrained("qquarkq/multitask-toxicity-classifier")

# Загрузка конфигурации с порогами
with open("config.json", "r") as f:
    config = json.load(f)
thresholds = config["thresholds"]

class MultiTaskToxicityEncoder(nn.Module):
    def __init__(self, model_name, dropout=0.2, freeze_encoder=False):
        super().__init__()

        self.encoder = AutoModel.from_pretrained(model_name)
        if freeze_encoder:
          for param in self.encoder.parameters():
              param.requires_grad = False
        self.hidden_size = self.encoder.config.hidden_size
        self.dropout = nn.Dropout(dropout)

        #Головы
        self.profanity_head = nn.Linear(self.hidden_size, 1)  # Нецензурная лексика
        self.threat_head = nn.Linear(self.hidden_size, 1)     # Угрозы
        self.illegal_head = nn.Linear(self.hidden_size, 1) # Незаконный контент



    def forward(self, input_ids, attention_mask):
        outputs = self.encoder(
            input_ids=input_ids,
            attention_mask=attention_mask,
            return_dict=True
        )
        cls_embedding = outputs.last_hidden_state[:, 0, :]
        cls_embedding = self.dropout(cls_embedding)

        profanity_logits = self.profanity_head(cls_embedding)
        threat_logits = self.threat_head(cls_embedding)
        illegal_logits = self.illegal_head(cls_embedding)

        return profanity_logits, threat_logits, illegal_logits

model = MultiTaskToxicityEncoder(model_name=config["model_name"])

def predict_toxicity(text):
    encoding = tokenizer(
        text,
        truncation=True,
        padding='max_length',
        max_length=128,
        return_tensors='pt'
    )
    
    with torch.no_grad():
        profanity_logits, threat_logits, illegal_logits = model(
            encoding['input_ids'], 
            encoding['attention_mask']
        )
    
    profanity_prob = torch.sigmoid(profanity_logits).item()
    threat_prob = torch.sigmoid(threat_logits).item()
    illegal_prob = torch.sigmoid(illegal_logits).item()

    profanity_pred = int(profanity_prob >= thresholds['profanity'])
    threat_pred = int(threat_prob >= thresholds['threat'])
    illegal_pred = int(illegal_prob >= thresholds['illegal'])
    return {
        'profanity': profanity_prob,
        'threat': threat_prob,
        'illegal': illegal_prob
    }

# Пример использования
text = "Ты полный идиот!"
result = predict_toxicity(text)
print(result)
````
Датасет
Модель обучена на датасете [qquarkq/russian-toxic-multilabel-comments](https://huggingface.co/datasets/qquarkq/russian-toxic-multilabel-comments)