---
language: ru
license: mit
tags:
- toxicity
- multi-task
- rubert-tiny2
- text-classification
metrics:
- f1
- precision
- recall
---

# Multi-Task Toxicity Classifier for Russian

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

- **Profanity** (мат/ненормативная лексика)
- **Threat** (угрозы)
- **Illegal** (запросы о незаконных действиях)

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

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

| Класс      | Порог | F1-score | Precision | Recall |
|------------|-------|----------|-----------|--------|
| Profanity  | 0.75 | 0.982 | 0.990 | 0.973 |
| Threat     | 0.50 | 1.000 | 1.000 | 1.000 |
| Illegal    | 0.10 | 0.997 | 1.000 | 0.994 |

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

```python
import torch
from transformers import AutoTokenizer
from huggingface_hub import hf_hub_download
import json

# Загрузка модели и токенизатора
model_path = hf_hub_download(repo_id="AlucardV/kinopotok-toxicity-multitask-model_update", filename="pytorch_model.bin")
config_path = hf_hub_download(repo_id="AlucardV/kinopotok-toxicity-multitask-model_update", filename="config.json")
tokenizer = AutoTokenizer.from_pretrained("AlucardV/kinopotok-toxicity-multitask-model_update")

# Определение класса модели
class MultiTaskToxicityEncoder(torch.nn.Module):
    def __init__(self, model_name):
        super().__init__()
        from transformers import AutoModel
        self.encoder = AutoModel.from_pretrained(model_name)
        hidden_size = self.encoder.config.hidden_size
        self.profanity_head = torch.nn.Linear(hidden_size, 1)
        self.threat_head = torch.nn.Linear(hidden_size, 1)
        self.illegal_head = torch.nn.Linear(hidden_size, 1)

    def forward(self, input_ids, attention_mask):
        outputs = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
        cls = outputs.last_hidden_state[:, 0, :]
        return {
            "profanity": self.profanity_head(cls).squeeze(-1),
            "threat": self.threat_head(cls).squeeze(-1),
            "illegal": self.illegal_head(cls).squeeze(-1),
        }

# Загрузка весов
config = json.load(open(config_path))
model = MultiTaskToxicityEncoder(config["model_name"])
model.load_state_dict(torch.load(model_path, map_location="cpu"))
model.eval()

# Функция предсказания
def predict(text):
    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=128)
    with torch.no_grad():
        outputs = model(inputs["input_ids"], inputs["attention_mask"])
    probs = {k: torch.sigmoid(v).item() for k, v in outputs.items()}
    thresholds = {
        "profanity": 0.7500000000000002,
        "threat": 0.5000000000000001,
        "illegal": 0.1,
    }
    results = {
        k: {
            "toxic": probs[k] >= thresholds[k],
            "confidence": f"{probs[k]*100:.1f}%"
        } for k in probs
    }
    return results

# Пример
print(predict("Ты что, совсем охренел, мудак?"))