Publish reproducible multi-task toxicity model

README.md

@@ -3,32 +3,101 @@ language: ru
 pipeline_tag: text-classification
 base_model: cointegrated/rubert-tiny2
 license: cc-by-nc-sa-4.0
+tags:
+- multi-label
+- toxicity
+- russian
+- custom-code
 ---
 
 # RuBERT Tiny 2 Multi-Task Toxicity
 
-Multi-label модель с общим замороженным энкодером и тремя независимыми
-линейными головами: profanity, threat, illegal.
+End-to-end multi-label encoder с общим RuBERT Tiny 2 и тремя
+независимыми линейными головами: `profanity`, `threat`, `illegal`.
+
+## Данные
+
+Нативные русские Toxic Russian Comments и RUCriminal после
+fail-closed удаления политических высказываний, политически
+чувствительных тем и явных персональных идентификаторов.
+Для `illegal` используются законные hard negatives по тем же
+статьям УК РФ. Переводные данные отсутствуют. Legal train/test
+разделены по непересекающимся статьям.
 
 ## Test metrics
 
-| class     |   support |   threshold |   precision |   recall |     f1 |   pr_auc |
-|:----------|----------:|------------:|------------:|---------:|-------:|---------:|
-| profanity |       181 |        0.5  |      0.6795 |   0.8785 | 0.7663 |   0.849  |
-| threat    |       155 |        0.64 |      0.6571 |   0.7419 | 0.697  |   0.7605 |
-| illegal   |        73 |        0.78 |      0.679  |   0.7534 | 0.7143 |   0.7813 |
+| class     |   support |   threshold |   precision |   recall |     f1 |   f1_ci_low |   f1_ci_high |   pr_auc |
+|:----------|----------:|------------:|------------:|---------:|-------:|------------:|-------------:|---------:|
+| profanity |       141 |        0.75 |      0.9452 |   0.9787 | 0.9617 |      0.9342 |       0.9817 |   0.9806 |
+| threat    |        91 |        0.83 |      0.8298 |   0.8571 | 0.8432 |      0.7858 |       0.8984 |   0.9167 |
+| illegal   |       103 |        0.5  |      0.9717 |   1      | 0.9856 |      0.9679 |       1      |   1      |
+
+Macro F1: 0.9302; Micro F1: 0.9369.
+
+Пороги выбраны только на validation:
+
+```json
+{
+  "profanity": 0.75,
+  "threat": 0.83,
+  "illegal": 0.5
+}
+```
+
+## Метрики по доменам
 
-Macro F1: 0.7258; Micro F1: 0.7319.
-Пороги выбраны только по validation.
+| domain          | class     |   support |     f1 |   recall |
+|:----------------|:----------|----------:|-------:|---------:|
+| legal_questions | illegal   |       102 | 1      |   1      |
+| social_comments | profanity |       141 | 0.9617 |   0.9787 |
+| social_comments | threat    |        91 | 0.8432 |   0.8571 |
+| social_comments | illegal   |         1 | 0.4    |   1      |
 
 ## Использование
 
-Архитектура MultiTaskToxicityEncoder и функция predict_toxicity(text)
-приведены в сопровождающем notebook. В репозитории сохраняются
-токенизатор, model_state.pt и model_config.json.
+```python
+import torch
+from transformers import (
+    AutoModelForSequenceClassification,
+    AutoTokenizer,
+)
+
+repo = "aurelianvolturi/rubert-tiny2-multitask-toxicity"
+tokenizer = AutoTokenizer.from_pretrained(repo)
+model = AutoModelForSequenceClassification.from_pretrained(
+    repo,
+    trust_remote_code=True,
+    low_cpu_mem_usage=False,
+    device_map=None,
+).eval()
+
+text = "Как проверить подлинность паспорта?"
+batch = tokenizer(
+    text,
+    return_tensors="pt",
+    truncation=True,
+    max_length=model.config.max_length,
+)
+
+with torch.inference_mode():
+    probabilities = torch.sigmoid(
+        model(**batch).logits
+    )[0].tolist()
+
+result = {
+    label: probability >= model.config.thresholds[label]
+    for label, probability in zip(
+        model.config.labels,
+        probabilities,
+    )
+}
+```
 
 ## Ограничения
 
-Модель выполняет предварительную модерацию, а не юридическую
-квалификацию. Редкие и завуалированные формулировки требуют
-ручной проверки.
+Модель предназначена для предварительной модерации и не выполняет
+юридическую квалификацию. Завуалированные формулировки, цитаты,
+ирония и новый сленг могут требовать ручной проверки.
+
+Custom code следует загружать только из доверенного репозитория и
+фиксировать `revision` на commit hash.

config.json

@@ -0,0 +1,81 @@
+{
+  "architectures": [
+    "MultiTaskToxicityEncoder"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_multitask_toxicity.MultiTaskToxicityConfig",
+    "AutoModelForSequenceClassification": "modeling_multitask_toxicity.MultiTaskToxicityEncoder"
+  },
+  "dropout": 0.15,
+  "dtype": "float32",
+  "encoder_config": {
+    "_name_or_path": "cointegrated/rubert-tiny2",
+    "add_cross_attention": false,
+    "architectures": [
+      "BertForPreTraining"
+    ],
+    "attention_probs_dropout_prob": 0.1,
+    "bos_token_id": null,
+    "chunk_size_feed_forward": 0,
+    "classifier_dropout": null,
+    "dtype": "float32",
+    "emb_size": 312,
+    "eos_token_id": null,
+    "gradient_checkpointing": false,
+    "hidden_act": "gelu",
+    "hidden_dropout_prob": 0.1,
+    "hidden_size": 312,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "initializer_range": 0.02,
+    "intermediate_size": 600,
+    "is_decoder": false,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "layer_norm_eps": 1e-12,
+    "max_position_embeddings": 2048,
+    "model_type": "bert",
+    "num_attention_heads": 12,
+    "num_hidden_layers": 3,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "pad_token_id": 0,
+    "position_embedding_type": "absolute",
+    "problem_type": null,
+    "return_dict": true,
+    "tie_word_embeddings": true,
+    "transformers_version": "5.13.1",
+    "type_vocab_size": 2,
+    "use_cache": true,
+    "vocab_size": 83828
+  },
+  "encoder_name": "cointegrated/rubert-tiny2",
+  "id2label": {
+    "0": "profanity",
+    "1": "threat",
+    "2": "illegal"
+  },
+  "label2id": {
+    "illegal": 2,
+    "profanity": 0,
+    "threat": 1
+  },
+  "labels": [
+    "profanity",
+    "threat",
+    "illegal"
+  ],
+  "max_length": 128,
+  "model_type": "multitask_toxicity",
+  "thresholds": {
+    "illegal": 0.5,
+    "profanity": 0.75,
+    "threat": 0.83
+  },
+  "transformers_version": "5.13.1"
+}

configuration_multitask_toxicity.py

@@ -0,0 +1,34 @@
+from transformers import PreTrainedConfig
+
+
+class MultiTaskToxicityConfig(PreTrainedConfig):
+    model_type = "multitask_toxicity"
+
+    def __init__(
+        self,
+        encoder_name="cointegrated/rubert-tiny2",
+        encoder_config=None,
+        labels=None,
+        thresholds=None,
+        max_length=128,
+        dropout=0.15,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.encoder_name = encoder_name
+        self.encoder_config = dict(encoder_config or {})
+        self.labels = list(
+            labels or ["profanity", "threat", "illegal"]
+        )
+        self.thresholds = thresholds or {
+            label: 0.5 for label in self.labels
+        }
+        self.max_length = int(max_length)
+        self.dropout = float(dropout)
+        self.num_labels = len(self.labels)
+        self.id2label = {
+            index: label for index, label in enumerate(self.labels)
+        }
+        self.label2id = {
+            label: index for index, label in enumerate(self.labels)
+        }

inference.py

@@ -0,0 +1,40 @@
+import torch
+from transformers import (
+    AutoModelForSequenceClassification,
+    AutoTokenizer,
+)
+
+REPO_ID = "aurelianvolturi/rubert-tiny2-multitask-toxicity"
+
+tokenizer = AutoTokenizer.from_pretrained(REPO_ID)
+model = AutoModelForSequenceClassification.from_pretrained(
+    REPO_ID,
+    trust_remote_code=True,
+    low_cpu_mem_usage=False,
+    device_map=None,
+).eval()
+
+
+def predict(text):
+    batch = tokenizer(
+        text,
+        truncation=True,
+        max_length=model.config.max_length,
+        return_tensors="pt",
+    )
+    with torch.inference_mode():
+        probabilities = torch.sigmoid(
+            model(**batch).logits
+        )[0].tolist()
+
+    return {
+        label: {
+            "detected":
+                probability >= model.config.thresholds[label],
+            "probability": probability,
+        }
+        for label, probability in zip(
+            model.config.labels,
+            probabilities,
+        )
+    }

model_state.pt → model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:d9756ef248659d3d7fba2ebf7bae2a66a0e054ccb70979e9d0bbf42be6de173e
-size 116802296
+oid sha256:379356c6a512c9cbca9687747dfa8e3a1cd0610f6e50a8494e9b44a36876561d
+size 116785916

model_config.json

@@ -1,15 +0,0 @@
-{
-  "encoder_name": "cointegrated/rubert-tiny2",
-  "labels": [
-    "profanity",
-    "threat",
-    "illegal"
-  ],
-  "thresholds": {
-    "profanity": 0.4999999999999998,
-    "threat": 0.6399999999999997,
-    "illegal": 0.7799999999999997
-  },
-  "max_length": 128,
-  "seed": 42
-}

modeling_multitask_toxicity.py

@@ -0,0 +1,62 @@
+import torch
+import torch.nn as nn
+from transformers import BertConfig, BertModel, PreTrainedModel
+from transformers.modeling_outputs import SequenceClassifierOutput
+
+from .configuration_multitask_toxicity import MultiTaskToxicityConfig
+
+
+class MultiTaskToxicityEncoder(PreTrainedModel):
+    config_class = MultiTaskToxicityConfig
+    base_model_prefix = "encoder"
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        if not config.encoder_config:
+            raise ValueError(
+                "encoder_config отсутствует в config.json"
+            )
+
+        # Не вызываем from_pretrained() внутри __init__. Внешний
+        # from_pretrained() сам создаёт модель и загружает все веса.
+        encoder_config = BertConfig.from_dict(
+            dict(config.encoder_config)
+        )
+        self.encoder = BertModel(encoder_config)
+        self.dropout = nn.Dropout(config.dropout)
+        self.heads = nn.ModuleDict({
+            label: nn.Linear(encoder_config.hidden_size, 1)
+            for label in config.labels
+        })
+        self.post_init()
+
+    def forward(
+        self,
+        input_ids,
+        attention_mask,
+        labels=None,
+        **kwargs,
+    ):
+        cls_embedding = self.encoder(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+        ).last_hidden_state[:, 0]
+
+        hidden = self.dropout(cls_embedding)
+        logits = torch.cat([
+            self.heads[label](hidden)
+            for label in self.config.labels
+        ], dim=1)
+
+        loss = None
+        if labels is not None:
+            loss = nn.functional.binary_cross_entropy_with_logits(
+                logits,
+                labels,
+            )
+
+        return SequenceClassifierOutput(
+            loss=loss,
+            logits=logits,
+        )