| ---
|
| 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("Ты что, совсем охренел, мудак?"))
|
| |
