基于Python的自然语言处理系列（42）：Token Classification（标注分类）

在本篇文章中，我们将探讨如何进行 Token Classification（标注分类），这是一类为句子中的每个 token（词或子词）分配标签的任务。该任务可以解决很多问题，例如命名实体识别（NER）、词性标注（POS）和分块（Chunking）。本文将聚焦于命名实体识别任务，并展示如何使用 BERT 模型进行微调。

1. 数据加载

我们使用 CoNLL-2003 数据集，这是一个常用的命名实体识别数据集。通过 load_dataset() 函数加载数据集：

python 复制代码

from datasets import load_dataset

raw_datasets = load_dataset("conll2003")

加载后的数据集包含三个任务的标签：NER、POS 和 Chunking。我们主要关注 NER 任务，因此接下来我们会处理 ner_tags 标签。

2. 数据预处理

在进行标注分类任务时，输入文本需要被转换为 token ID。由于我们处理的是预分词数据（即每个输入已按词分割），我们需要使用 is_split_into_words=True 参数告诉 tokenizer 如何处理这些分词数据。首先，下载并缓存 BERT 预训练模型的 tokenizer：

python 复制代码

from transformers import AutoTokenizer

model_checkpoint = "bert-base-cased"
tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)

我们可以用 tokenizer 处理预分词输入，同时利用 word_ids() 方法确保 token 与标签正确对齐：

python 复制代码

inputs = tokenizer(raw_datasets["train"][0]["tokens"], is_split_into_words=True)
inputs.word_ids()

通过处理后的输入，特殊 token（如 [CLS] 和 [SEP]）被添加到序列中，此外模型对一些词进行了子词切分。我们还需要确保标签能够正确匹配这些子词，特别是处理切分后的子词标签对齐问题。

3. 标签对齐

接下来，我们编写一个函数，用于将标签与 tokenizer 生成的 token 对齐。特殊 token 将被标记为 -100（这在损失计算中会被忽略），而子词的标签会与该词的第一个 token 一致。具体代码如下：

python 复制代码

def align_labels_with_tokens(labels, word_ids):
    new_labels = []
    current_word = None
    for word_id in word_ids:
        if word_id != current_word:
            current_word = word_id
            label = -100 if word_id is None else labels[word_id]
            new_labels.append(label)
        elif word_id is None:
            new_labels.append(-100)
        else:
            label = labels[word_id]
            if label % 2 == 1:
                label += 1
            new_labels.append(label)
    return new_labels

然后我们可以对整个数据集进行预处理，将 tokenizer 和标签对齐函数应用到整个数据集上：

python 复制代码

def tokenize_and_align_labels(examples):
    tokenized_inputs = tokenizer(examples["tokens"], truncation=True, is_split_into_words=True)
    all_labels = examples["ner_tags"]
    new_labels = []
    for i, labels in enumerate(all_labels):
        word_ids = tokenized_inputs.word_ids(i)
        new_labels.append(align_labels_with_tokens(labels, word_ids))
    tokenized_inputs["labels"] = new_labels
    return tokenized_inputs

tokenized_datasets = raw_datasets.map(tokenize_and_align_labels, batched=True, remove_columns=raw_datasets["train"].column_names)

4. 模型准备

我们将使用 AutoModelForTokenClassification 类加载 BERT 模型。首先，需要指定模型的标签映射：

python 复制代码

id2label = {str(i): label for i, label in enumerate(label_names)}
label2id = {v: k for k, v in id2label.items()}

然后加载模型：

python 复制代码

from transformers import AutoModelForTokenClassification

model = AutoModelForTokenClassification.from_pretrained(model_checkpoint, id2label=id2label, label2id=label2id)

5. 模型训练与评估

接下来我们创建训练和评估所需的 DataLoader，并定义损失函数和优化器。使用 DataCollatorForTokenClassification 来确保输入和标签在填充时保持相同的长度：

python 复制代码

from transformers import DataCollatorForTokenClassification
from torch.utils.data import DataLoader

data_collator = DataCollatorForTokenClassification(tokenizer=tokenizer)

train_dataloader = DataLoader(tokenized_datasets["train"], shuffle=True, collate_fn=data_collator, batch_size=8)
eval_dataloader = DataLoader(tokenized_datasets["validation"], collate_fn=data_collator, batch_size=8)

from torch.optim import AdamW
optimizer = AdamW(model.parameters(), lr=2e-5)

接下来，使用 accelerator 和学习率调度器进行训练，并保存模型到 Hugging Face Hub：

python 复制代码

from accelerate import Accelerator
from transformers import get_scheduler

accelerator = Accelerator()
model, optimizer, train_dataloader, eval_dataloader = accelerator.prepare(model, optimizer, train_dataloader, eval_dataloader)

num_train_epochs = 3
num_training_steps = num_train_epochs * len(train_dataloader)
lr_scheduler = get_scheduler("linear", optimizer=optimizer, num_warmup_steps=0, num_training_steps=num_training_steps)

# 开始训练
for epoch in range(num_train_epochs):
    model.train()
    for batch in train_dataloader:
        outputs = model(**batch)
        loss = outputs.loss
        accelerator.backward(loss)
        optimizer.step()
        lr_scheduler.step()
        optimizer.zero_grad()

6. 推理

模型训练完成后，我们可以使用 Hugging Face 的 pipeline 进行推理：

python 复制代码

from transformers import pipeline

model_checkpoint = "Chaklam/bert-finetuned-ner-accelerate"
token_classifier = pipeline("token-classification", model=model_checkpoint, aggregation_strategy="simple")

print(token_classifier("My name is Peter and I work at AIT in Bangkok, Thailand."))

通过这些步骤，你将能够微调并部署一个命名实体识别模型。

结语

在本篇文章中，我们介绍了如何使用 BERT 模型进行 Token Classification（标注分类）任务。我们从数据的加载与预处理开始，详细讲解了如何处理预分词数据，并将标签与模型的输入 tokens 对齐。接着，我们通过 BERT 模型进行微调，使用 DataLoader 来构建训练与验证数据集，最终完成了模型的训练与评估。在最后，我们展示了如何通过 Hugging Face 的 pipeline 进行命名实体识别的推理。

标注分类任务（如命名实体识别）在自然语言处理中有广泛的应用，本篇文章展示的 BERT 模型微调流程可以为多种标注任务提供基础。无论是命名实体识别、词性标注还是分块识别，都可以使用类似的方法进行模型微调与部署。通过对该过程的深入理解，你将能够在实际项目中应用这些技术，解决更多复杂的 NLP 问题。

如果你觉得这篇博文对你有帮助，请点赞、收藏、关注我，并且可以打赏支持我！

欢迎关注我的后续博文，我将分享更多关于人工智能、自然语言处理和计算机视觉的精彩内容。

谢谢大家的支持！