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

1. 数据加载

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

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:

from transformers import AutoTokenizer

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

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

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

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

3. 标签对齐

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

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 和标签对齐函数应用到整个数据集上:

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 模型。首先,需要指定模型的标签映射:

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

        然后加载模型:

from transformers import AutoModelForTokenClassification

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

5. 模型训练与评估

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

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:

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 进行推理:

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 问题。

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

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

谢谢大家的支持!

点赞(0) 打赏

评论列表 共有 0 条评论

暂无评论

微信公众账号

微信扫一扫加关注

发表
评论
返回
顶部