资讯动态

Transformers 实战指南:使用 DistilBERT 微调 Token 分类(NER 命名实体识别)模型

发布时间:2026/9/10 11:28:31 来源:尧图企业网站定制
Transformers 实战指南使用 DistilBERT 微调 Token 分类NER 命名实体识别模型【免费下载链接】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/transformersToken 分类Token Classification是自然语言处理中的一类基础任务它要求模型为输入序列中的每个 token预测一个标签。本文基于本仓库Hugging Face Transformers在 docs/source/ja/tasks/token_classification.md 中提供的完整指南以最经典的**命名实体识别NER**为例手把手带你完成从数据集加载、标签对齐预处理、seqeval 指标评估到基于Trainer微调 DistilBERT再到用pipeline与原生 PyTorch 两种方式进行推理的端到端流程。读完本文你将能够独立复现一个可检测人名、地名、组织名等实体的 NER 模型并理解其背后的源码级实现原理。任务背景什么是 Token 分类与 NERToken 分类的任务定义很直接给定一句话为其中的每一个 token分配一个类别标签。最典型的应用就是命名实体识别NER——找出句中的人Person、地点Location、组织Organization等实体。本文的实战目标是在 WNUT 17 数据集上微调 DistilBERT使其能够检测出训练数据中未见过的新实体然后把微调好的模型投入实际推理。理解 BIO 标注体系在开始前先理解 NER 数据集的标注体系。以 WNUT 17 数据集的一条训练样本为例 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中的每个数字都对应一个实体标签 ID。把数字 ID 转换为可读的标签名 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, ]这就是经典的BIOBegin / Inside / Outside标注体系每个ner_tag前缀的含义如下B-Begin表示该 token 是一个实体的开始I-Inside表示该 token属于同一个实体内部。例如Empire State Building这个实体中State和Building都带I-前缀0Outside表示该 token不属于任何实体。从上面的示例可以看到Empire标签 7 B-location、State、Building标签 8 I-location共同构成了一个完整的地点实体而ESB又被单独标记为B-location。环境准备与登录开始之前请确保已安装所需库pip install transformers datasets evaluate seqeval各库的职责如下库作用transformers提供预训练模型、Tokenizer、Trainer、Pipeline 等核心组件datasets加载和处理 WNUT 17 等数据集evaluate加载并计算 seqeval 等评估指标seqeval专为序列标注NER设计的指标库可计算精确率、召回率、F1、准确率另外建议登录 Hugging Face 账号以便把训练好的模型上传到 Hub 与社区共享。在提示符出现时输入你的 access token 即可 from huggingface_hub import notebook_login notebook_login()提示本文所有涉及 Transformer 架构与检查点的兼容性说明都可以在本仓库的 token 分类任务文档体系中找到对应章节如需查看当前支持 Token 分类任务的完整架构列表可参考任务页面的说明。加载 WNUT 17 数据集使用datasets库的load_dataset一行即可加载 WNUT 17 from datasets import load_dataset wnut load_dataset(wnut_17)数据集的划分split包含train训练集、validation验证集与test测试集其中train和test会被用于后续微调与评估。预处理加载 Tokenizer 并处理词—子词错位加载 Tokenizer 与is_split_into_words加载 DistilBERT 的 tokenizer from transformers import AutoTokenizer tokenizer AutoTokenizer.from_pretrained(distilbert/distilbert-base-uncased)虽然上面的样本看起来已经分词了但那是按空格切分的单词并没有经过 BERT 系模型的WordPiece 子词切分。为了让 tokenizer 直接把单词列表当作输入、并在此基础上进一步切成子词必须显式传入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]]注意观察两个关键变化序列首尾被自动加上了特殊 token[CLS]和[SEP]部分单词被拆成了子词例如paulwalk→paul##walkESB→es##b。这就产生了输入与标签之间的错位原来一个词对应一个标签现在一个词可能对应多个子词 token。因此必须重新对齐 token 与标签。标签重对齐的三条核心规则对齐逻辑由三条规则构成这也是整个 Token 分类预处理中最关键的部分使用word_ids方法把所有 token 映射回它所属的原始单词BatchEncoding.word_ids返回与input_ids等长的列表每个位置记录该 token 对应第几个原始单词特殊 token 对应None给特殊 token[CLS]、[SEP]分配标签-100——这是 PyTorchCrossEntropyLoss约定俗成的忽略索引loss 计算时会自动跳过这些位置对应源码 modeling_distilbert.py 中loss_fct CrossEntropyLoss(); loss loss_fct(logits.view(-1, self.num_labels), labels.view(-1))的实现只给每个单词的第一个子词 token 打标签同一单词的其余子词也赋-100。编写对齐函数并应用将上述规则封装成预处理函数 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) # 将 token 映射回其所属单词 ... previous_word_idx None ... label_ids [] ... for word_idx in word_ids: # 特殊 token 置为 -100 ... if word_idx is None: ... label_ids.append(-100) ... elif word_idx ! previous_word_idx: # 只给每个单词的第一个 token 打标签 ... 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这里truncationTrue会把超过模型最大输入长度DistilBERT 为 512的序列截断。然后对整个数据集批量应用该函数——batchedTrue表示一次处理多条样本以加速 tokenized_wnut wnut.map(tokenize_and_align_labels, batchedTrue)用 DataCollatorForTokenClassification 动态填充接下来需要把样本组装成 batch。与其把整个数据集都 pad 到最大长度更高效的做法是在组 batch 时动态填充到当前 batch 内最长序列的长度。Transformers 为此专门提供了DataCollatorForTokenClassification from transformers import DataCollatorForTokenClassification data_collator DataCollatorForTokenClassification(tokenizertokenizer)该 collator 的源码位于 data_collator.py它做两件关键的事对input_ids、attention_mask按 batch 内最长长度动态 padpaddingTrue即longest策略对 labels 同步填充右侧 padding 时在标签序列末尾补上label_pad_token_id默认-100从而保证模型计算 loss 时不会把 padding 位置算进去。该 collator 还支持max_length指定最大长度、pad_to_multiple_of将长度 pad 到某数值的倍数可配合 NVIDIA Volta 及以上架构的 Tensor Core 使用等参数以及return_tensorspt返回 PyTorch 张量或npNumPy 数组。评估指标用 seqeval 计算 NER 分数训练过程中引入评估指标能直观反映模型性能。本任务使用evaluate库加载seqeval指标——它是序列标注任务的行业标准指标库能同时输出**精确率precision、召回率recall、F1 和准确率accuracy**等多个分数 import evaluate seqeval evaluate.load(seqeval)先取出真实标签对应的标签名再编写compute_metrics函数把模型的 logits 取 argmax 得到预测类别 ID过滤掉-100的位置后将 ID 映射回标签名交给seqeval.compute计算整体指标 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], ... }说明因为-100位置对应特殊 token 与被忽略的子词在评估前必须把它们从预测和标签中一并过滤掉否则会干扰分数计算——这正是上面两个列表推导式里if l ! -100的作用。compute_metrics函数会在后续配置Trainer时传入。训练配置标签映射与 Trainer建立 id2label / label2id 映射开始训练前先建立标签 ID 与标签名之间的双向映射。id2label用于把模型输出的类别 ID 映射为标签名label2id用于训练时的标签编码 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, ... }加载 Token 分类头使用AutoModelForTokenClassification加载 DistilBERT并通过num_labels13指定类别数同时传入两个映射 from transformers import AutoModelForTokenClassification, TrainingArguments, Trainer model AutoModelForTokenClassification.from_pretrained( ... distilbert/distilbert-base-uncased, num_labels13, id2labelid2label, label2idlabel2id ... )从源码看AutoModelForTokenClassification会自动路由到对应的架构实现。以本任务的 DistilBERT 为例其 Token 分类实现位于 modeling_distilbert.py模型主体DistilBertModel之上叠加一个nn.Dropout和一层nn.Linear(config.hidden_size, config.num_labels)作为分类头对每个位置的隐层输出做线性映射得到形状为(batch_size, sequence_length, num_labels)的 logitsforward 中若传入labels则把 logits 和 labels 都展开成(-1,)后交给CrossEntropyLoss计算 loss形状为(batch_size, sequence_length)的labels展开后长度与 logits 对齐-100位置自动被忽略。可见加载时传num_labels会自动替换输出分类头正是通过这种结构实现的。配置 TrainingArguments 并启动训练剩余步骤只有三步在TrainingArguments中定义训练超参数。唯一必填参数是模型保存目录output_dir设置push_to_hubTrue可以把模型推送到 Hub需先登录 Hugging Face在每个 epoch 结束时Trainer会评估指标并保存训练 checkpoint把训练参数连同模型、数据集、tokenizer、数据 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()这里各超参数的典型含义与建议取值参数取值说明output_dirmy_awesome_wnut_model模型与 checkpoint 的保存目录必填learning_rate2e-5微调预训练模型常用的较小学习率per_device_train_batch_size16每张卡设备上的训练 batch 大小per_device_eval_batch_size16每张卡上的评估 batch 大小num_train_epochs2训练轮数weight_decay0.01AdamW 优化器的权重衰减系数帮助抑制过拟合eval_strategyepoch每个 epoch 结束时评估一次save_strategyepoch每个 epoch 结束时保存 checkpointload_best_model_at_endTrue训练结束后自动加载评估指标最优的 checkpointpush_to_hubTrue训练结束后把模型推送到 Hub需登录训练完成后用push_to_hub把模型分享出去方便社区复用 trainer.push_to_hub()提示如果你还不太熟悉用Trainer微调模型建议先阅读本文所依据任务文档中关于 Train with PyTorch Trainer 的基础教程再回到这里更完整的 Token 分类微调示例还可以参考仓库中 PyTorch 方向的 token-classification 示例examples/pytorch/token-classification 目录下的run_ner.py。推理两种方式使用微调模型微调完成后模型即可投入推理。以下以一段 NBA 球队描述文本为例 text The Golden State Warriors are an American professional basketball team based in San Francisco.方式一使用 pipeline推荐最简单的方式是把微调模型包装进pipeline。用任务标识符ner实例化分类器直接把文本丢进去 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、token 在序列中的序号index、还原后的原始单词word以及该单词在原文本中的起止字符位置start/end——利用这两个位置可以很方便地在原句中高亮实体。pipeline背后对应的是TokenClassificationPipeline实现位于 pipelines/token_classification.py它内部会完成 tokenize、前向推理、argmax 取类别、按id2label映射回标签名等一系列步骤。它还支持一些实用参数例如aggregation_strategy默认为none不聚合按 token 输出设为simple或first/max/average时会把B-/I-属于同一实体的连续 token聚合成一个实体输出entity_group字段更贴近实际使用场景ignore_labels指定忽略哪些标签例如忽略O默认忽略O使输出只保留实体。方式二手动复现 pipeline原生 PyTorch如果你需要更多控制权也可以手动复现 pipeline 的推理流程一共三步。第一步用微调模型的 tokenizer 对文本分词并返回 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注意这里用torch.no_grad()关闭梯度计算推理时无需反向传播可以节省显存并加速。第三步取 logits 在类别维dim2上的最大值下标再通过模型的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]对照输入文本逐 token 查看Golden State被标为B-locationI-location地点Warriors被标为B-group组织San与Francisco被标为B-locationB-location注意这里两个词各被预测为实体开头属于模型输出中的边界不完美情况也是 NER 推理中的常见现象。id2label之所以可用正是因为训练时我们把它传给了AutoModelForTokenClassification.from_pretrained它会写入模型配置推理时即可直接从model.config.id2label读取。小结本文完整走通了基于 Transformers 的 Token 分类NER微调全流程从理解 BIO 标注体系、加载 WNUT 17 数据集到用is_split_into_wordsword_ids解决词—子词标签错位、用DataCollatorForTokenClassification动态 padding 与-100标签填充再到用 seqeval 评估、Trainer微调 DistilBERT最后通过pipeline或原生 PyTorch 完成推理。回顾几个最值得记住的要点标签对齐是 Token 分类预处理的核心难点-100既是CrossEntropyLoss的忽略标记也是 collator 的默认标签填充值它贯穿了预处理 → 训练 → 评估过滤整条链路num_labelsid2label/label2id决定了分类头的结构与可解释性它们会随模型一起保存供推理时直接读取评估时必须过滤-100位置否则 seqeval 的分数会被特殊 token 与子词位置污染。掌握了这套流程你只需替换数据集、标签体系和基础模型就能把同样的方法迁移到词性标注POS、分块Chunking等其他序列标注任务上。【免费下载链接】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 小时内与您沟通定制方案

免费获取报价