用 DistilBERT 做一个内容审核 AI 模型

YWD2023 · 2026-04-17 21:56:40 · 科技·工程

需要检测用户输入文本的场景很多，比如说（当你要做）：评论区、论坛、聊天室……等等。

如果直接不做任何处理，平台做大了各种广告显然就会肆无忌惮地刷屏，影响用户体验，某些内容还可能是违法违规的。被找上门责令改正并罚款就完蛋了。

其实这个文章的主要写作目的是我自己想做聊天室系统，要检测内容……

显然，你没那个能力去做人工审核（用户发个帖子等几十分钟才能审核完发出去？）需要自动判断哪些正常、哪些不正常。

正则？关键词匹配？

简单场景应该也许可能能用，但文本有变化就不行了。规则是写不完的，维护词表的难度也很大。

而且有的词汇单独出现时可能无害，但和其他词组合在一起就成了其他意思，通杀就太严格了，放任不管又太宽松了。

AI 模型能不能会这些复杂的关系？

所以我打算弄一个模型，能自己从标注数据里学规律。

经过和 DeepSeek 的一番友好选型交流，最后我选用了 HuggingFace 的 DistilBERT 预训练模型来搭建审核 AI。

DistilBERT 的内部结构

Transformer 层与自注意力

DistilBERT 的核心是 Transformer 架构。

Transformer 之所以能理解输入的文本，依赖的是自注意力机制。

自注意力让每个 token 在编码时都能看到序列里的所有其他 token，而不是像那些 RNN 那样逐字顺序处理。

具体来说，给定一个 token 序列，自注意力把每个 token 的向量映射成三个向量：查询（Q）、键（K）、值（V）。然后对于 token i，它的输出是序列中所有 token 的 V 的加权和，权重由 Q 和 K 的点积决定：

其中，d_k 是向量维度；除以 \sqrt{d_k} 是为了防止点积过大，导致 \operatorname{softmax} 的梯度过小。权重向量经过 \operatorname{softmax} 后会变成概率分布，表示“我在理解当前 token 时应该把注意力放在哪些位置上”。

DistilBERT 使用多头注意力（Multi-Head Attention），也就是说 DistilBERT 把上面的过程并行做 12 次，每个头用不同的 W_Q、W_K、W_V 投影矩阵，学习不同类型的关注关系（有的头关注句法，有的头关注语义）。12 个头的输出拼接后再做一次线性变换，得到最终的注意力输出。

每个 Transformer 层的完整结构是：

也就是：

残差连接保证梯度能顺畅地反向传播穿过多层网络。

DistilBERT

满血原版 BERT-base 高达有 12 个 Transformer 层，隐层维度为 768，总参数约为 1.1 \times 10^8。

DistilBERT 把层数砍到 6 层，维度保持为 768。参数量从 1.1 \times 10^8 降到 6.6 \times 10^7，比直接用 BERT 快不少，毕竟家没矿没那么多算力。

BERT 在每条输入序列的开头插入一个特殊的 [CLS]（classification）token。这个 token 没有实际词义，但经过 6 层 Transformer 后，它的隐状态向量会通过注意力机制聚合整个序列的语义信息。分类任务就是取这个 [CLS] 位置的 768 维向量，送进分类头。

（另一个特殊 token 是 [SEP]，用于分隔两段文本，现在的内容识别只有一段在句尾。）

项目

我用了 distilbert-base-multilingual-cased 模型，支持 104 种语言（包括中文）。

源代码可在 GitHub 获取。train.py 是训练脚本，predict.py 是运行脚本，模型文件可以在 GitHub Releases 下载。

数据我选择了 Olivi-9/Chinese-Spam-Dataset 进行测试，其中的格式是两列 CSV：

label,text
0,"正常文本"
1,"广告文本"

其中，0 表示正常，1 表示敏感。

后续，我又选择了 Chinese-Offensive-Language-Dataset 中的 SexHarmSet 和 AbuseSet 数据集并与前面的数据集合并进行测试，其中的格式是三列 CSV：

Keyword,Type,Sentence
XX,Safe,正常内容
XX,Harmful,敏感内容

为了方便训练，后面我把这两种格式的 CSV 转成统一的 label,text 格式（丢弃了 Chinese-Offensive-Language-Dataset 中的 Keyword 字段），并合并成一个新的数据集。

环境准备

pip install torch transformers pandas scikit-learn accelerate

如果有 CUDA GPU，确保先装了对应版本的 PyTorch（去 pytorch.org 按你的 CUDA 版本选）不然默认安装只能用 CPU。

CPU 也能训练就是极慢无比。

训练流程

1. 加载和切分数据

import pandas as pd
from sklearn.model_selection import train_test_split

df = pd.read_csv("train.csv")
train_texts, val_texts, train_labels, val_labels = train_test_split(
    df["text"].tolist(),
    df["label"].tolist(),
    test_size=0.2,
    random_state=42,
    stratify=df["label"],  # 保持正负比例
)

stratify 主要是防止正负样本比例不均衡，随机切分可能导致集里某一类占比太多/太少力。加了 stratify 后每个子集都保持原始比例。

对于数据不均衡的情况，见下文的处理数据部分。

2. 分词编码

包括 DistilBERT 在内的 Transformer 模型（是的，AIer 爱用的 LLM 也是 Transformer 模型）都不能直接处理文本，需要先用 Tokenizer 把文字转成 token ID 序列。

这里用的分词算法是 WordPiece，是一种子词切分方案。

先解释一下为什么要用子词分词。如果词表是整词级别的（像中文词那样），碰到生僻词、变形词就直接变成未知。

如果词表是字符级别的，词表很小，但序列会很长，理解更难。

WordPiece 把高频词保留为完整单元，把低频词拆成更小的子词片段，用 ## 前缀标记“这是上一个词的延续”。

例如：

from transformers import DistilBertTokenizerFast

tokenizer = DistilBertTokenizerFast.from_pretrained("distilbert-base-multilingual-cased")

enc = tokenizer("来加我微信，加我可以领取 EXAMPLE 优惠券", return_tensors="pt")

print(enc["input_ids"])
# tensor([[  101,  4501,  2598,  3976,  3785,  2316, 10064,  2598,  3976,  2756, 2204,  8427,  2737, 80426, 36535, 11127, 51036,  2224,  3887,  2558,  102]])

print(tokenizer.convert_ids_to_tokens(enc["input_ids"][0]))
# ['[CLS]', '来', '加', '我', '微', '信', '，', '加', '我', '可', '以', '领', '取', 'EX', '##AM', '##P', '##LE', '优', '惠', '券', '[SEP]']

序列开头是 [CLS]（ID 为 101），结尾是 [SEP]（ID 为 102）。这两个特殊 token 是模型要求的格式，训练和推理都必须有。

中文字符基本上每个字是一个 token（多语言 BERT 模型对中文基本上是字级别切分），到了英文是 subword 级别。

批量 tokenize 时，padding=True 会把同一 batch 里所有序列 pad 到相同长度，并生成 attention_mask，其中 1 表示真实 token，0 表示 padding 位置。模型在做注意力计算时，padding 位置的注意力分数会被设为 -\infty；经过 \operatorname{softmax} 后，对应权重会趋近于 0，不参与实际计算。

from transformers import DistilBertTokenizerFast

MODEL_NAME = "distilbert-base-multilingual-cased"
tokenizer = DistilBertTokenizerFast.from_pretrained(MODEL_NAME)

train_encodings = tokenizer(train_texts, truncation=True, padding=True, max_length=256)
val_encodings = tokenizer(val_texts, truncation=True, padding=True, max_length=256)

max_length=256：广告文本通常不长，256 个 token 足够覆盖。如果你的文本特别长可以调大，但不能超过 512，512 是 DistilBERT 位置编码的硬上限（位置编码只预训练了 512 个位置，超过就没学过）。

3. 包装成 Dataset

PyTorch 的训练循环靠 DataLoader 来管理数据：自动打乱顺序、按批大小切批次、多进程预加载。DataLoader 要求传入一个实现了 __getitem__(idx) 和 __len__() 的 Dataset 对象。前者按索引返回单条样本，后者告诉它数据集有多大。HuggingFace 的 Trainer 内部正是用 DataLoader 消耗数据的，所以我们要把编码结果套进这个格式。

__getitem__ 里做了一件关键的事：torch.tensor(val[idx]) 把 tokenizer 输出的 Python 列表（input_ids[i]、attention_mask[i] 等）即时转成 PyTorch tensor。不在这里转的原因是把整个数据集一次性转成 tensor 会大量占内存，按需转换更合理。

import torch
from torch.utils.data import Dataset

class TextDataset(Dataset):
    def __init__(self, encodings, labels):
        self.encodings = encodings
        self.labels = labels

    def __getitem__(self, idx):
        item = {key: torch.tensor(val[idx]) for key, val in self.encodings.items()}
        item["labels"] = torch.tensor(self.labels[idx])
        return item

    def __len__(self):
        return len(self.labels)

train_dataset = TextDataset(train_encodings, train_labels)
val_dataset = TextDataset(val_encodings, val_labels)

DataLoader 的 collate_fn（批次拼接函数）负责把 __getitem__ 返回的多个样本字典拼成一个批次字典，input_ids 的形状会从 \left[\text{seq\_len}\right] 变成 \left[\text{batch\_size}, \text{seq\_len}\right]，其他字段同理。Trainer 使用默认的 collate_fn，已经支持了这个格式，不需要我们手动写。

4. 前向传播

一条样本进入模型：

Embedding

input_ids（整数序列）通过嵌入矩阵查表，变成 token 向量序列，每个 token 映射到 768 维的稠密向量。同时加上位置编码，同样学来的 768 维向量，告诉模型每个 token 在序列中的位置信息（因为自注意力本身不感知顺序，位置信息必须显式注入）。两者相加：

6 层 Transformer