Rubert Tiny Toxicity

🚀 微调版俄文文本毒性分类模型

本项目基于 cointegrated/rubert-tiny 模型微调，用于对俄语短文本（如社交网络评论）的毒性和不当性进行分类。

标签信息

快速测试

你可以通过以下示例文本快速测试模型：

"Иди ты нафиг!"

🚀 快速开始

这是 cointegrated/rubert-tiny 模型的微调版本，用于对俄语短非正式文本（如社交网络中的评论）的毒性和不当性进行分类。

问题被表述为多标签分类，包含以下类别：

• non-toxic：根据 OK ML Cup 竞赛的定义，文本不包含侮辱、脏话和威胁。
• insult：侮辱
• obscenity：脏话
• threat：威胁
• dangerous：根据 Babakov 等人 的定义，文本不合适，即可能损害发言者的声誉。

如果文本同时为 non-toxic 且非 dangerous，则可认为该文本是安全的。

💻 使用示例

基础用法

以下函数用于估计文本具有毒性或危险性的概率：

# !pip install transformers sentencepiece --quiet
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification

model_checkpoint = 'cointegrated/rubert-tiny-toxicity'
tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
model = AutoModelForSequenceClassification.from_pretrained(model_checkpoint)
if torch.cuda.is_available():
    model.cuda()
    
def text2toxicity(text, aggregate=True):
    """ Calculate toxicity of a text (if aggregate=True) or a vector of toxicity aspects (if aggregate=False)"""
    with torch.no_grad():
        inputs = tokenizer(text, return_tensors='pt', truncation=True, padding=True).to(model.device)
        proba = torch.sigmoid(model(**inputs).logits).cpu().numpy()
    if isinstance(text, str):
        proba = proba[0]
    if aggregate:
        return 1 - proba.T[0] * (1 - proba.T[-1])
    return proba

print(text2toxicity('я люблю нигеров', True))
# 0.9350118728093193

print(text2toxicity('я люблю нигеров', False))
# [0.9715758  0.0180863  0.0045551  0.00189755 0.9331106 ]

print(text2toxicity(['я люблю нигеров', 'я люблю африканцев'], True))
# [0.93501186 0.04156357]

print(text2toxicity(['я люблю нигеров', 'я люблю африканцев'], False))
# [[9.7157580e-01 1.8086294e-02 4.5550885e-03 1.8975559e-03 9.3311059e-01]
#  [9.9979788e-01 1.9048342e-04 1.5297388e-04 1.7452303e-04 4.1369814e-02]]

🔧 技术细节

该模型在 OK ML Cup 和 Babakov 等人 的联合数据集上进行训练，使用 Adam 优化器，学习率为 1e-5，批量大小为 64，训练了 15 个周期，具体训练过程可参考 Colab 笔记本。

如果文本的不当性得分高于 0.8，则认为该文本不合适；如果得分低于 0.2，则认为该文本合适。开发集上每个标签的 ROC AUC 如下：

non-toxic  : 0.9937
insult     : 0.9912
obscenity  : 0.9881
threat     : 0.9910
dangerous  : 0.8295