| --- |
| language: |
| - ru |
| license: mit |
| pipeline_tag: text-classification |
| tags: |
| - toxicity |
| - multi-task |
| - rubert |
| - moderation |
| metrics: |
| - f1 |
| --- |
| |
| # 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) |
|
|
| ```python |
| 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) |
| |