--- language: ru license: mit tags: - toxicity - multilabel-classification - russian - rubert-tiny2 - multitask-learning pretty_name: RuBERT Tiny2 Toxicity Multitask model_name: rubert-tiny2-toxicity-multitask --- # Model Card for rubert-tiny2-toxicity-multitask ## Model Details **Model Description:** Эта модель представляет собой многозадачный классификатор токсичности для русского языка. Она построена на основе лёгкого энкодера `cointegrated/rubert-tiny2` и содержит три независимые классификационные «головы» для одновременного предсказания трёх классов: - **profanity** – нецензурная лексика и оскорбления - **threat** – угрозы - **illegal** – запросы, связанные с незаконными действиями Модель обучена в мультитаск-режиме (Multi-Task Learning) с использованием взвешенной функции потерь для учёта дисбаланса классов. **Разработчик:** Arrtemwolf **Модель доступна на Hugging Face:** [https://huggingface.co/Arrtemwolf/rubert-tiny2-toxicity-multitask](https://huggingface.co/Arrtemwolf/rubert-tiny2-toxicity-multitask) ## Intended Use **Назначение:** Модель предназначена для автоматической модерации текстовых сообщений в социальных сетях, чатах, комментариях. Она может использоваться для: - фильтрации нецензурной брани; - выделения угроз для приоритетной проверки модераторами; - обнаружения сообщений, содержащих запросы на совершение незаконных действий. **Ограничения:** - Модель обучена на ограниченном наборе данных (около 5500 примеров) и может не охватывать все разнообразие русскоязычной токсичности. - Рекомендуется использовать в связке с дополнительными фильтрами и человеческой модерацией для критичных применений. ## Training Data Модель обучена на датасете, собранном из двух открытых источников: - [petaevd/russian-toxicity-multilabel](https://huggingface.co/datasets/petaevd/russian-toxicity-multilabel) – для класса `illegal` - [AlexSham/Toxic_Russian_Comments](https://huggingface.co/datasets/AlexSham/Toxic_Russian_Comments) – для классов `profanity` и `threat` Датасет был очищен, сбалансирован и опубликован отдельно: [Arrtemwolf/ru-toxic-multi-label-clean](https://huggingface.co/datasets/Arrtemwolf/ru-toxic-multi-label-clean) ## Evaluation Results Модель оценивалась на валидационной выборке (20% от общего датасета). Для каждого класса был подобран оптимальный порог вероятности, максимизирующий **F1-меру**. | Класс | Порог | Precision | Recall | F1-score | |------------|-------|-----------|--------|----------| | profanity | 0.50 | 0.821 | 0.908 | **0.862**| | threat | 0.15 | 0.261 | 0.571 | **0.358**| | illegal | 0.40 | 1.000 | 1.000 | **1.000**| > Примечание: высокий F1 для класса `illegal` обусловлен небольшим количеством положительных примеров в валидации; на более крупных выборках результаты могут отличаться. ## How to Use ### Загрузка модели и токенизатора ```python from transformers import AutoTokenizer, AutoModel import torch MODEL_NAME = "Arrtemwolf/rubert-tiny2-toxicity-multitask" tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME) encoder = AutoModel.from_pretrained(MODEL_NAME) # Загрузка обученных голов (классификаторов) # Внимание: головы сохранены отдельно, их нужно загрузить и прикрепить к модели # Ниже приведён пример класса-обёртки, который можно использовать после загрузки весов голов. Для удобства рекомендуется использовать класс MultiTaskToxicityEncoder, который объединяет энкодер и три головы. Веса голов сохранены в файле multitask_heads.pt в репозитории. Пример загрузки: python class MultiTaskToxicityEncoder(torch.nn.Module): def __init__(self, encoder): super().__init__() self.encoder = encoder hidden_size = encoder.config.hidden_size self.head_profanity = torch.nn.Linear(hidden_size, 1) self.head_threat = torch.nn.Linear(hidden_size, 1) self.head_illegal = torch.nn.Linear(hidden_size, 1) self.dropout = torch.nn.Dropout(0.3) 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, :] cls_embedding = self.dropout(cls_embedding) return (self.head_profanity(cls_embedding), self.head_threat(cls_embedding), self.head_illegal(cls_embedding)) # Загружаем энкодер encoder = AutoModel.from_pretrained(MODEL_NAME) model = MultiTaskToxicityEncoder(encoder) # Загружаем веса голов state_dict = torch.load("multitask_heads.pt", map_location="cpu") model.load_state_dict(state_dict, strict=False) # strict=False, т.к. веса только для голов model.eval() Предсказание для одного текста python def predict(text, model, tokenizer, device="cpu"): encoded = tokenizer(text, padding=True, truncation=True, max_length=256, return_tensors="pt") input_ids = encoded["input_ids"].to(device) attention_mask = encoded["attention_mask"].to(device) with torch.no_grad(): logit_p, logit_t, logit_i = model(input_ids, attention_mask) prob_p = torch.sigmoid(logit_p).item() prob_t = torch.sigmoid(logit_t).item() prob_i = torch.sigmoid(logit_i).item() # Пороги (оптимальные, полученные на валидации) thresholds = {"profanity": 0.50, "threat": 0.15, "illegal": 0.40} return { "profanity": {"prob": prob_p, "label": prob_p >= thresholds["profanity"]}, "threat": {"prob": prob_t, "label": prob_t >= thresholds["threat"]}, "illegal": {"prob": prob_i, "label": prob_i >= thresholds["illegal"]}, } # Пример text = "Ты мне угрожаешь? Я вызову полицию!" result = predict(text, model, tokenizer) print(result)