Software 2.0

微调 BERT 实现命名实体识别

admin -

命名实体识别（NER）是自然语言处理（NLP）中的一项基本任务，涉及识别和分类文本中的命名实体到预定义的类别，如人名、组织、地点等。NER 有广泛的应用，从改进搜索引擎结果到为聊天机器人提供动力和从文档中提取关键信息。

在这篇技术博客文章中，我们将通过一个完整的示例来使用流行的 CoNLL-2003 数据集微调预训练的 BERT 模型进行命名实体识别。我们将使用 Hugging Face transformers 和 datasets 库来构建我们的 NER 模型。

1、设置环境

首先，我们需要安装必要的 Python 库。我们将使用 transformers 作为模型和分词器，datasets 来加载数据，evaluate 和 seqeval 用于性能指标，以及其他有用的库。

!pip install -q accelerate>=0.27.2 peft>=0.9.0 bitsandbytes>=0.43.0 transformers>=4.38.2 trl>=0.7.11 sentencepiece>=0.1.99
!pip install -q sentence-transformers>=3.0.0 mteb>=1.1.2 datasets>=2.18.0

在我们的环境设置好之后，让我们导入必要的模块。

from transformers import AutoModelForTokenClassification, AutoTokenizer
from transformers import DataCollatorWithPadding
from transformers import TrainingArguments, Trainer
import numpy as np
from datasets import load_dataset, Dataset

2、加载和准备数据

我们将使用 CoNLL-2003 数据集，这是一个 NER 的基准数据集。它包含标注了四种实体类型的新闻文章：人员（PER）、组织（ORG）、地点（LOC）和杂项（MISC）。我们可以从 Hugging Face Hub 轻松加载它。

# The CoNLL-2003 dataset for NER
dataset = load_dataset("conll2003", trust_remote_code=True)

让我们从训练集中查看一个示例，以了解数据结构。

example = dataset["train"][848]
example

这将输出一个包含 tokens、词性（POS）标签、chunk 标签以及最重要的 NER 标签的字典。

{'id': '848',
 'tokens': ['Dean',
  'Palmer',
  'hit',
  'his',
  '30th',
  'homer',
  'for',
  'the',
  'Rangers',
  '.'],
 'pos_tags': [22, 22, 38, 29, 16, 21, 15, 12, 23, 7],
 'chunk_tags': [11, 12, 21, 11, 12, 12, 13, 11, 12, 0],
 'ner_tags': [1, 2, 0, 0, 0, 0, 0, 0, 3, 0]}

ner_tags 由整数表示。为了使它们更易读，我们将创建从 ID 到标签的映射，反之亦然。CoNLL-2003 数据集使用 IOB（内部、外部、开始）格式，其中 B- 表示实体的开始，I- 表示实体内部的标记，O 表示标记不是实体的一部分。

label2id = {
    'O': 0, 'B-PER': 1, 'I-PER': 2, 'B-ORG': 3, 'I-ORG': 4,
    'B-LOC': 5, 'I-LOC': 6, 'B-MISC': 7, 'I-MISC': 8
}
id2label = {index: label for label, index in label2id.items()}

3、模型和分词器

我们将使用预训练的 bert-base-cased 模型。这个模型是许多 NLP 任务的良好起点，包括 NER。我们将使用 transformers 库中的 AutoTokenizer 和 AutoModelForTokenClassification 加载分词器和模型。

from transformers import AutoModelForTokenClassification

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

# Load model
model = AutoModelForTokenClassification.from_pretrained(
    "bert-base-cased",
    num_labels=len(id2label),
    id2label=id2label,
    label2id=label2id
)

准备数据的一个关键步骤是将标签与分词器生成的标记对齐。BERT 的分词器经常将单词分解为子标记。例如，单词"Rangers"可能会被分词为"Range"和"##rs"。我们需要确保我们的标签正确映射到这些子标记。

以下函数，align_labels，负责处理这个问题。它遍历标记及其相应的标签，对于每个标记，它确定其子标记的正确标签。特殊标记如 [CLS] 和 [SEP] 被分配标签 -100，这是让损失函数在训练期间忽略它们的标准做法。

def align_labels(examples):
    token_ids = tokenizer(examples["tokens"], truncation=True, is_split_into_words=True)
    labels = examples["ner_tags"]

    updated_labels = []
    for index, label in enumerate(labels):
        word_ids = token_ids.word_ids(batch_index=index)
        previous_word_idx = None
        label_ids = []
        for word_idx in word_ids:
            if word_idx != previous_word_idx:
                previous_word_idx = word_idx
                updated_label = -100 if word_idx is None else label[word_idx]
                label_ids.append(updated_label)
            elif word_idx is None:
                label_ids.append(-100)
            else:
                updated_label = label[word_idx]
                if updated_label % 2 == 1:
                    updated_label += 1
                label_ids.append(updated_label)

        updated_labels.append(label_ids)

    token_ids["labels"] = updated_labels
    return token_ids

tokenized_dataset = dataset.map(align_labels, batched=True)

4、评估

对于评估我们的 NER 模型，我们将使用 seqeval 库，这是序列标记任务的标准。我们将计算 F1 分数，它是精确率和召回率的调和平均值，并提供了模型性能的更平衡的度量。

下面的 compute_metrics 函数将由 Trainer 用于在训练期间评估模型。

import evaluate

seqeval = evaluate.load("seqeval")

def compute_metrics(eval_pred):
    logits, labels = eval_pred
    predictions = np.argmax(logits, axis=2)

    true_predictions = []
    true_labels = []

    for prediction, label in zip(predictions, labels):
      for token_prediction, token_label in zip(prediction, label):
        if token_label != -100:
          true_predictions.append([id2label[token_prediction]])
          true_labels.append([id2label[token_label]])

    results = seqeval.compute(predictions=true_predictions, references=true_labels)
    return {"f1": results["overall_f1"]}

5、训练模型

现在，我们准备设置训练过程。我们将使用 transformers 库中的 Trainer 类，它简化了训练循环。我们定义训练参数，包括学习率、批大小和训练轮数。

from transformers import DataCollatorForTokenClassification

data_collator = DataCollatorForTokenClassification(tokenizer=tokenizer)

training_args = TrainingArguments(
   "model",
   learning_rate=2e-5,
   per_device_train_batch_size=16,
   per_device_eval_batch_size=16,
   num_train_epochs=1,
   weight_decay=0.01,
   save_strategy="epoch",
   report_to="none"
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset["train"],
    eval_dataset=tokenized_dataset["test"],
    tokenizer=tokenizer,
    data_collator=data_collator,
    compute_metrics=compute_metrics,
)

trainer.train()

训练完成后，我们可以在测试集上评估微调后的模型。

trainer.evaluate()

The output will show the evaluation loss and the F1-score on the test data. 
An F1-score around 0.90 indicates that our model is performing well.

{'eval_loss': 0.13917773962020874,
 'eval_f1': 0.9067902220802044,
 'eval_runtime': 11.171,
 'eval_samples_per_second': 309.103,
 'eval_steps_per_second': 19.336,
 'epoch': 1.0}

6、推理

最后，让我们使用微调后的模型对新句子进行预测。我们可以使用 transformers 中的 pipeline 函数进行简单的推理。

from transformers import pipeline

trainer.save_model("ner_model")

token_classifier = pipeline(
    "token-classification",
    model="ner_model",
)
token_classifier("My name is Damen.I want to play in NBA")

输出将是一个字典列表，其中每个字典代表一个识别出的实体、其标签及其在文本中的位置。

[{'entity': 'B-PER',
  'score': np.float32(0.9887255),
  'index': 4,
  'word': 'Dame',
  'start': 11,
  'end': 15},
 {'entity': 'I-PER',
  'score': np.float32(0.9901882),
  'index': 5,
  'word': '##n',
  'start': 15,
  'end': 16},
 {'entity': 'B-MISC',
  'score': np.float32(0.77015245),
  'index': 12,
  'word': 'NBA',
  'start': 35,
  'end': 38}]

如您所见，模型正确地将"Damen"识别为人员（PER），将"NBA"识别为杂项实体（MISC）。

7、结束语

在这篇博客文章中，我们介绍了微调预训练 BERT 模型进行命名实体识别的过程。我们涵盖了从加载和准备数据到训练、评估和使用模型进行推理的所有内容。Hugging Face 生态系统使得只需几行代码就可以构建强大的 NLP 模型变得异常简单。您可以将此工作流程适应您自己的 NER 任务和数据集，以构建自定义实体识别系统。

原文链接: Fine-Tuning BERT for Named Entity Recognition: A Step-by-Step Guide

汇智网翻译整理，转载请标明出处