资讯动态

使用 Transformers 微调 DistilBERT 完成 Token Classification(词元分类)实战指南

发布时间:2026/9/9 23:53:48 来源:尧图企业网站定制
使用 Transformers 微调 DistilBERT 完成 Token Classification词元分类实战指南【免费下载链接】transformers Transformers: the model-definition framework for state-of-the-art machine learning models in text, vision, audio, and multimodal models, for both inference and training.项目地址: https://gitcode.com/GitHub_Trending/tra/transformers导读词元分类Token Classification是自然语言处理的基础任务之一它要求模型为句子中的每一个词元token预测一个标签。本文以命名实体识别NER为典型场景基于当前transformers仓库带你从零开始将 DistilBERT 在 WNUT 17 数据集上微调成可识别人、地点、组织等实体的专用模型并给出完整的预处理、训练配置、评估与推理代码。读完本文你将掌握词元与标签的对齐技巧、DataCollatorForTokenClassification的填充原理、基于 seqeval 的 NER 评估体系以及pipeline与手动推理两种落地方式。任务定义为什么词元分类需要逐词元建模Token classification assigns a label to individual tokens in a sentence词元分类为句子中的每个词元分配一个标签。在 NER 中模型要对每个词元回答它是不是一个命名实体、是什么类型的实体。这一任务与整句分类text-classification最大的区别在于输出粒度模型对每个位置都要产出一个分类结果。在transformers仓库中该任务由以*ForTokenClassification结尾的模型类承担例如 DistilBertForTokenClassification。从其实现可以看到任务头classification head的构造方式在编码器输出的sequence_output之上加一层 Dropout 与一层nn.Linear(config.hidden_size, config.num_labels)线性分类器把每个位置的隐藏向量映射为num_labels维的 logitsself.distilbert DistilBertModel(config) self.dropout nn.Dropout(config.dropout) self.classifier nn.Linear(config.hidden_size, config.num_labels)训练时若传入labels则把 logits 展平为(-1, num_labels)后与同样展平的标签计算CrossEntropyLossmodeling_distilbert.py#L794-L797。这也是后文用-100屏蔽特殊词元与子词这一技巧的底层原因CrossEntropyLoss会默认忽略目标值为-100的位置。从源码的自动映射表 MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING_NAMES 可以看出支持该任务的架构覆盖面极广BERT、DistilBERT、RoBERTa、ELECTRA、DeBERTa、GPT-2、Llama、Mistral、Gemma、Falcon 等数十种模型都实现了对应的分类头。因此本文以 DistilBERT 为例讲解的完整流水线可无缝迁移到上述任一架构。环境准备与数据集加载先安装本次实战所需的全部依赖pip install transformers datasets evaluate seqeval其中seqeval是 NER/词元分类场景的标准评估库它基于 BIOBegin/Inside/Outside标注体系计算精确率、召回率、F1 与准确率。提示本文使用TrainerAPI 微调。若你尚未接触过该 API可先阅读 trainer 使用指南。训练完成后还可以把模型上传并分享到社区需先在终端按提示输入访问令牌完成登录 from huggingface_hub import notebook_login notebook_login()加载 WNUT 17 数据集WNUT 17 是一个专为稀有/未知实体识别设计的评测数据集。使用datasets库一行即可加载 from datasets import load_dataset wnut load_dataset(wnut_17)查看一条训练样本 wnut[train][0] {id: 0, ner_tags: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 7, 8, 8, 0, 7, 0, 0, 0, 0, 0, 0, 0, 0], tokens: [paulwalk, It, s, the, view, from, where, I, m, living, for, two, weeks, ., Empire, State, Building, , ESB, ., Pretty, bad, storm, here, last, evening, .]}样本由两个字段组成tokens是已经切好的词序列ner_tags是与每个词一一对应的标签编号。把编号还原为标签名即可看清标注内容 label_list wnut[train].features[fner_tags].feature.names label_list [ O, B-corporation, I-corporation, B-creative-work, I-creative-work, B-group, I-group, B-location, I-location, B-person, I-person, B-product, I-product, ]WNUT 17 共定义 6 类实体corporation机构、creative-work作品、group群体、location地点、person人名、product产品每一类都遵循 BIO 标注。读懂 BIO 标注约定每个标签前的字母代表词元在实体中的位置这是理解整个任务与评估逻辑的关键B-Begin标注一个实体的起始词例如B-location表示一个地点实体的开头I-Inside标注同一实体内部的后续词。例如EmpireB-location、StateI-location、BuildingI-location共同组成多词地点Empire State BuildingOOutside表示该词不属于任何实体。B-前缀还有一个语义作用当两个同类实体相邻时如两个连续的地点B-用于标示第二个实体的重新开始这正是基于B-/I-前缀做序列解码与评估的前提。预处理子词切分后的标签重对齐加载 DistilBERT 分词器 from transformers import AutoTokenizer tokenizer AutoTokenizer.from_pretrained(distilbert/distilbert-base-uncased)数据集中的tokens字段看起来已经tokenized但实际上只是按空白/词法完成的词切分。要让模型能处理还需把每个词进一步拆成子词subword。由于tokens是词列表而非原始字符串需要在调用分词器时指定is_split_into_wordsTrue example wnut[train][0] tokenized_input tokenizer(example[tokens], is_split_into_wordsTrue) tokens tokenizer.convert_ids_to_tokens(tokenized_input[input_ids]) tokens [[CLS], , paul, ##walk, it, , s, the, view, from, where, i, , m, living, for, two, weeks, ., empire, state, building, , es, ##b, ., pretty, bad, storm, here, last, evening, ., [SEP]]对照可见两个关键变化序列首尾被加入特殊词元[CLS]与[SEP]出现了子词拆分例如paulwalk被拆成paul##walkESB被拆成es##b。这直接导致输入序列与标签在长度和位置上失配原来一个词对应一个标签拆分后可能对应多个子词。因此必须在预处理阶段把标签重新对齐到子词序列上规则如下使用分词器的word_ids()方法建立每个 token 属于哪个原始词的映射将特殊词元[CLS]、[SEP]word_ids返回None的标签设为-100使其在CrossEntropyLoss中被自动忽略对同一原始词拆分出的多个子词只保留第一个子词的标签其余子词同样设为-100。基于以上规则实现tokenize_and_align_labels并用truncationTrue把超长样本截断到模型的最大输入长度 def tokenize_and_align_labels(examples): ... tokenized_inputs tokenizer(examples[tokens], truncationTrue, is_split_into_wordsTrue) ... labels [] ... for i, label in enumerate(examples[fner_tags]): ... word_ids tokenized_inputs.word_ids(batch_indexi) # Map tokens to their respective word. ... previous_word_idx None ... label_ids [] ... for word_idx in word_ids: # Set the special tokens to -100. ... if word_idx is None: ... label_ids.append(-100) ... elif word_idx ! previous_word_idx: # Only label the first token of a given word. ... label_ids.append(label[word_idx]) ... else: ... label_ids.append(-100) ... previous_word_idx word_idx ... labels.append(label_ids) ... tokenized_inputs[labels] labels ... return tokenized_inputs应用map处理整个数据集batchedTrue能利用底层批处理加速 tokenized_wnut wnut.map(tokenize_and_align_labels, batchedTrue)提示word_ids()只在 fast tokenizer基于 Rust tokenizers 的PreTrainedTokenizerFast上可用。这也是选择 fast 版分词器默认加载的理由之一。DataCollatorForTokenClassification动态填充Trainer在每步训练时会从数据集中抽出一个 batch。由于样本长度不一需要用 Data Collator 把同一个 batch 内的样本补齐到相同长度。推荐使用DataCollatorForTokenClassification做动态填充——只把当前 batch 内序列填充到最长长度而不是把全数据集填充到统一的最大长度从而显著减少无效计算 from transformers import DataCollatorForTokenClassification data_collator DataCollatorForTokenClassification(tokenizertokenizer)从该类的源码data/data_collator.py可以看到它的核心参数与行为参数默认值说明tokenizer必填用于按模型的 padding 侧与 padding index 填充paddingTrueTrue/longest填充到 batch 内最长序列max_length填充到max_lengthFalse不填充max_lengthNone与paddingmax_length配合使用的上限长度pad_to_multiple_ofNone填充到该值的整数倍利于在 Volta 及以后的 NVIDIA GPU 上启用 Tensor Coreslabel_pad_token_id-100标签填充所用的占位 id-100会被 PyTorch 损失函数自动忽略return_tensorspt返回张量类型可选pt或np实现上它把labels从样本中抽出对input_ids、attention_mask等做动态 padding再按tokenizer.padding_side把labels用label_pad_token_id填充到与输入等长并转为int64张量。这样即使trainer传入的样本标签长度不一也能与输入严格对齐并进入损失计算。评估指标用 seqeval 计算 NER 分数用evaluate库加载seqeval。它不仅能输出准确率还会输出 NER 领域更关键的精确率、召回率与 F1 import evaluate seqeval evaluate.load(seqeval)compute_metrics函数是Trainer的钩子会在每个评估点接收模型的原始predictions形状为[batch, seq_len, num_labels]的 logits与labels。函数需要完成两步先用argmax取每个位置概率最高的类别 id再剔除标签为-100的位置这些是特殊词元与被屏蔽的子词不参与打分最后把 id 还原为标签名交给 seqeval import numpy as np labels [label_list[i] for i in example[fner_tags]] def compute_metrics(p): ... predictions, labels p ... predictions np.argmax(predictions, axis2) ... true_predictions [ ... [label_list[p] for (p, l) in zip(prediction, label) if l ! -100] ... for prediction, label in zip(predictions, labels) ... ] ... true_labels [ ... [label_list[l] for (p, l) in zip(prediction, label) if l ! -100] ... for prediction, label in zip(predictions, labels) ... ] ... results seqeval.compute(predictionstrue_predictions, referencestrue_labels) ... return { ... precision: results[overall_precision], ... recall: results[overall_recall], ... f1: results[overall_f1], ... accuracy: results[overall_accuracy], ... }注意compute_metrics的入参p是(predictions, labels)的元组这里的labels变量是解包出来的真实标签张量与上方示例行的labels列表无关。seqeval 对标签序列按实体级别而非 token 级别匹配因此它能正确惩罚把B-location, I-location断开成两个实体之类的边界错误。训练加载模型并配置 Trainer准备 id2label / label2id 映射模型输出的是类别 id为了让结果可读且推理时能自动反查标签名需要建立双向映射并注入模型配置 id2label { ... 0: O, ... 1: B-corporation, ... 2: I-corporation, ... 3: B-creative-work, ... 4: I-creative-work, ... 5: B-group, ... 6: I-group, ... 7: B-location, ... 8: I-location, ... 9: B-person, ... 10: I-person, ... 11: B-product, ... 12: I-product, ... } label2id { ... O: 0, ... B-corporation: 1, ... I-corporation: 2, ... B-creative-work: 3, ... I-creative-work: 4, ... B-group: 5, ... I-group: 6, ... B-location: 7, ... I-location: 8, ... B-person: 9, ... I-person: 10, ... B-product: 11, ... I-product: 12, ... }加载带分类头的预训练模型用AutoModelForTokenClassification加载 DistilBERT并通过num_labels指定输出类别数13 类同时传入两个映射。库会自动把模型的顶层分类头替换为 13 类的随机初始化线性层只保留编码器的预训练权重 from transformers import AutoModelForTokenClassification, TrainingArguments, Trainer model AutoModelForTokenClassification.from_pretrained( ... distilbert/distilbert-base-uncased, num_labels13, id2labelid2label, label2idlabel2id ... )配置 TrainingArguments 并启动训练随后只需三步在TrainingArguments中定义训练超参。唯一必填参数是output_dir模型保存目录。若想训练结束后把模型推送到 Hub 社区可设push_to_hubTrue需已登录。eval_strategyepoch让 Trainer 在每个 epoch 结束时用你的compute_metrics评测并打印 seqeval 分数同时保存 checkpoint把训练参数连同模型、数据集、tokenizer以processing_class传入、data collator 与compute_metrics一起交给Trainer调用trainer.train()开始微调。 training_args TrainingArguments( ... output_dirmy_awesome_wnut_model, ... learning_rate2e-5, ... per_device_train_batch_size16, ... per_device_eval_batch_size16, ... num_train_epochs2, ... weight_decay0.01, ... eval_strategyepoch, ... save_strategyepoch, ... load_best_model_at_endTrue, ... push_to_hubTrue, ... ) trainer Trainer( ... modelmodel, ... argstraining_args, ... train_datasettokenized_wnut[train], ... eval_datasettokenized_wnut[test], ... processing_classtokenizer, ... data_collatordata_collator, ... compute_metricscompute_metrics, ... ) trainer.train()上述超参学习率2e-5、batch size16、2 个 epoch、权重衰减0.01是该类任务经过验证的通用起点。训练结束后调用trainer.push_to_hub()即可把最终模型与config.json内含id2label/label2id一起上传供社区直接使用 trainer.push_to_hub()推理两种落地方式方式一使用 pipeline推荐pipeline把分词、前向、标签反查、与原文对齐等步骤全部封装好是快速体验与线上服务的最简路径。用ner任务标识符加载微调后的模型pipeline 源码实现 text The Golden State Warriors are an American professional basketball team based in San Francisco. from transformers import pipeline classifier pipeline(ner, modelstevhliu/my_awesome_wnut_model) classifier(text) [{entity: B-location, score: 0.42658573, index: 2, word: golden, start: 4, end: 10}, {entity: I-location, score: 0.35856336, index: 3, word: state, start: 11, end: 16}, {entity: B-group, score: 0.3064001, index: 4, word: warriors, start: 17, end: 25}, {entity: B-location, score: 0.65523505, index: 13, word: san, start: 80, end: 83}, {entity: B-location, score: 0.4668663, index: 14, word: francisco, start: 84, end: 93}]输出的每一项都包含entity标签名、score置信度、word文本片段以及start/end该片段在原始字符串中的字符偏移可直接用于高亮展示。这里的model参数可替换为你自己output_dir中保存的本地模型或训练后推送的任意仓库名。从 TokenClassificationPipeline 的源码可看到该 pipeline 还暴露了若干进阶参数ignore_labels默认[O]从输出中过滤掉哪些标签默认剔除O后仅保留实体aggregation_strategy实体聚合策略取值为none/simple/first/average/max定义见 AggregationStrategy。simple会按 BIO 规则把(B-TAG)…(I-TAG)连续片段合并为一个实体输出若模型把多词实体误切成多个B-可用first/max/average结合子词信息按词合并。注意基于词合并first/max/average要求使用 fast tokenizerstride滑窗处理超长文本同样只支持 fast tokenizerstride文本超过model_max_length时按步长切块处理重叠 token 数即stride要求aggregation_strategy非none。方式二手动复现 pipeline 内部流程不依赖pipeline时可以自己复现其三步流程这在需要把模型嵌入自定义服务或调试时非常有用。第一步分词并返回 PyTorch 张量 from transformers import AutoTokenizer tokenizer AutoTokenizer.from_pretrained(stevhliu/my_awesome_wnut_model) inputs tokenizer(text, return_tensorspt)第二步前向传播取得logits from transformers import AutoModelForTokenClassification model AutoModelForTokenClassification.from_pretrained(stevhliu/my_awesome_wnut_model) with torch.no_grad(): ... logits model(**inputs).logits第三步在最后一个维度上取argmax得到每个位置的预测类别 id再通过模型的config.id2label正是训练时注入的映射翻译成标签文本 predictions torch.argmax(logits, dim2) predicted_token_class [model.config.id2label[t.item()] for t in predictions[0]] predicted_token_class [O, O, B-location, I-location, B-group, O, O, O, O, O, O, O, O, B-location, B-location, O, O]可见模型成功识别出Golden StateB-location/I-location、WarriorsB-group、San Francisco两个B-location等实体。手动方式的输出是逐词元粒度的标签序列未做实体聚合与去重因此San与Francisco会被标注为两个独立的B-location这是与pipeline默认输出格式的差异所在——若需合并需要自行实现或改用pipeline的聚合策略。小结与延伸本文以 WNUT 17 上的命名实体识别为例完整走通了数据集加载 → BIO 标签理解 → 子词重对齐 → 动态填充 → seqeval 评估 → Trainer 微调 → 双路径推理的 Token Classification 标准流程。核心要点可归结为三条标签对齐决定训练成败子词切分会破坏词与标签的一一对应必须基于word_ids()只给每个原始词的第一个子词打标签并把特殊词元与其余子词标记为-100-100贯穿始终它在预处理中屏蔽不参与学习的词元在DataCollatorForTokenClassification中作为label_pad_token_id兜底 padding并被compute_metrics用作过滤点评估与推理需区分粒度seqeval 按实体B/I 边界评估pipeline可按聚合策略输出实体级结果而手动推理得到的是 token 级标签。该流水线不局限于 NER——词性标注POS、中文分词、方面级情感分析等同样属于 Token Classification只要把数据集换成对应标注如 UPOS/XPOS 标签并调整label_list即可复用。若需要完整的可运行训练脚本与更多工程化细节仓库中还提供了 pytorch 文本分类示例、任务相关的 测试用例 以及 TokenClassificationPipeline 的端到端行为定义可作进一步参考。【免费下载链接】transformers Transformers: the model-definition framework for state-of-the-art machine learning models in text, vision, audio, and multimodal models, for both inference and training.项目地址: https://gitcode.com/GitHub_Trending/tra/transformers创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考

读完文章,也想定制专属网站?

尧图设计师 24 小时内与您沟通定制方案

免费获取报价