SafeTechDev/Russian-Spam-classifier

AIMS-RU-Binary-Spam-Classifier

Бинарный классификатор спама для русскоязычных сообщений, часть проекта AIMS (AI Moderation System).

Модель дообучена поверх DeepPavlov/rubert-base-cased и определяет, является ли сообщение SPAM или SAFE. Лёгкая и быстрая модель.

Метрики

Оценка проводилась на внешнем датасете alt-gnome/telegram-spam, не пересекающемся с обучающей выборкой:

Метрика	Значение
F1	0.94
Precision	0.97
Recall	0.91

Высокий precision (0.97) означает низкую долю ложных срабатываний — модель редко помечает безопасные сообщения как спам, что важно для автоматической модерации без риска задеть живых пользователей.

Архитектура

• Энкодер: DeepPavlov/rubert-base-cased
• Голова классификации: Dropout(0.2) → Linear(hidden_size, 1)
• Пулинг: эмбеддинг токена [CLS]
• Функция потерь при обучении: BCEWithLogitsLoss с весом положительного класса (pos_weight) для компенсации дисбаланса SAFE/SPAM
• Выход: один логит → sigmoid → вероятность класса SPAM

РЕКОМЕНДАЦИЯ Нормализация текста

Перед подачей в модель рекомендуется нормализовать текст, это использовалось при обучении.

• убирает "разрежённое" написание (п.р.и.в.е.т. → привет)
• заменяет гомоглифы (похожие символы из других алфавитов — ᴀ, α, латиница вместо кириллицы и т.д.) на кириллические эквиваленты
• удаляет ссылки и команды ботов (/cmd@bot)
• схлопывает повторяющиеся буквы (приввееееет → привет)
• приводит текст к нижнему регистру

Быстрый запуск

pip install torch transformers huggingface_hub

import re
import json
import torch
import torch.nn as nn
from transformers import AutoTokenizer, AutoModel
from huggingface_hub import hf_hub_download

REPO = "SafeTechDev/AIMS-RU-Binary-Spam-Classifier"
device = "cuda" if torch.cuda.is_available() else "cpu"

# ── Архитектура ──────────────────────────────────────────────────────────
class BinaryModel(nn.Module):
    def __init__(self, model_name):
        super().__init__()
        self.bert = AutoModel.from_pretrained(model_name, low_cpu_mem_usage=True)
        hidden = self.bert.config.hidden_size
        self.binary = nn.Sequential(
            nn.Dropout(0.2),
            nn.Linear(hidden, 1)
        )

    def forward(self, input_ids, attention_mask):
        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        pooled = outputs.last_hidden_state[:, 0]
        return self.binary(pooled).squeeze(-1)

# ── Загрузка ─────────────────────────────────────────────────────────────
config_path = hf_hub_download(REPO, "config.json")
weights_path = hf_hub_download(REPO, "pytorch_model.bin")

with open(config_path, encoding="utf-8") as f:
    cfg = json.load(f)

MAX_LENGTH = cfg.get("max_length", 40)

tokenizer = AutoTokenizer.from_pretrained(REPO)

model = BinaryModel(cfg["model"]).to(device)
model.load_state_dict(torch.load(weights_path, map_location=device, weights_only=True))
model.eval()

# ── Инференс ─────────────────────────────────────────────────────────────
def classify(text: str) -> dict:
    enc = tokenizer(
        text, truncation=True, padding="max_length",
        max_length=MAX_LENGTH, return_tensors="pt"
    )
    with torch.no_grad():
        logits = model(enc["input_ids"].to(device), enc["attention_mask"].to(device))

    prob_spam = float(torch.sigmoid(logits).squeeze().cpu().item())
    label = "SPAM" if prob_spam >= 0.7 else "SAFE"

    return {"label": label, "prob_spam": prob_spam}


print(classify("Дам денег, работу, пишите в лс"))
# {'label': 'SPAM', 'prob_spam': 0.98...}

print(classify("Привет, как дела?"))
# {'label': 'SAFE', 'prob_spam': 0.02...}

Через pipeline (упрощённо, без кастомной нормализации)

Модель хранит веса в формате pytorch_model.bin с кастомной головой, поэтому напрямую через transformers.pipeline("text-classification", ...) она не запустится — используйте код выше.

Файлы в репозитории

Файл	Назначение
pytorch_model.bin	Веса модели (state_dict для BinaryModel)
config.json	model (базовый энкодер), max_length
tokenizer.json, vocab.txt, ...	Файлы токенизатора rubert-base-cased

Ограничения

• Модель обучена и валидирована на русскоязычном тексте; на других языках качество не гарантируется.
• Порог 0.7 используется по умолчанию для разделения SAFE/SPAM — в проде рекомендуется подбирать порог под свои данные (например, по максимальному F1), т.к. соотношение precision/recall меняется в зависимости от домена.
• Модель не учитывает контекст переписки — классифицируется только отдельное сообщение.

Автор

SafeTechDev