Transformers documentation

Token 分类

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v5.22.2).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

Token 分类

Token classification 为句子中的单个 token 分配标签。最常见的 token 分类任务之一是命名实体识别(NER)。NER 旨在为句子中的每个实体(如人物、地点或组织)找到对应的标签。

本指南将向您展示如何:

  1. 在 WNUT 17 数据集上微调 DistilBERT,以检测新实体。
  2. 使用微调后的模型进行推理。

要查看与此任务兼容的所有架构和检查点,我们建议查阅任务页面。

在开始之前,请确保已安装所有必要的库:

pip install transformers datasets evaluate seqeval

我们建议您登录 Hugging Face 账户,以便上传模型并与社区共享。当系统提示时,请输入您的令牌以完成登录:

>>> from huggingface_hub import notebook_login

>>> notebook_login()

加载 WNUT 17 数据集

首先,从 🤗 Datasets 库中加载 WNUT 17 数据集:

>>> from datasets import load_dataset

>>> wnut = load_dataset("flaitenberger/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', '.']
}

ner_tags 中的每个数字代表一个实体。将这些数字转换为对应的标签名称,即可了解这些实体是什么:

>>> label_list = wnut["train"].features[f"ner_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",
]

每个 ner_tag 前缀的字母表示该实体中 token 的位置:

  • B- 表示实体的开始。
  • I- 表示该 token 包含在同一个实体内部(例如,State token 是 Empire State Building 这类实体的一部分)。
  • 0 表示该 token 不对应任何实体。

预处理

下一步是加载 DistilBERT 分词器,以预处理 tokens 字段:

>>> from transformers import AutoTokenizer

>>> tokenizer = AutoTokenizer.from_pretrained("distilbert/distilbert-base-uncased")

正如您在上面的示例 tokens 字段中所见,输入似乎已经被分词了。但实际上输入尚未分词,您需要设置 is_split_into_words=True 才能将单词分词为子词。例如:

>>> example = wnut["train"][0]
>>> tokenized_input = tokenizer(example["tokens"], is_split_into_words=True)
>>> 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],并且子词分词会导致输入与标签之间出现错位。原本对应单个标签的单个单词现在可能被拆分为两个子词。你需要通过以下方式重新对齐标记和标签:

  1. 使用 word_ids 方法将所有 token 映射到其对应的单词。
  2. 将标签 -100 分配给特殊标记 [CLS] 和 [SEP],以便 PyTorch 损失函数(参见 CrossEntropyLoss)忽略它们。
  3. 仅对给定单词的第一个 token 进行标注。将 -100 分配给同一单词的其他子 token。

以下是如何创建一个函数来重新对齐 token 和标签,并将序列截断为不超过 DistilBERT 的最大输入长度:

>>> def tokenize_and_align_labels(examples):
...     tokenized_inputs = tokenizer(examples["tokens"], truncation=True, is_split_into_words=True)

...     labels = []
...     for i, label in enumerate(examples[f"ner_tags"]):
...         word_ids = tokenized_inputs.word_ids(batch_index=i)  # 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

要对整个数据集应用预处理函数,请使用 🤗 Datasets 的 map 函数。你可以通过设置 batched=True 来一次处理数据集中的多个元素,从而加速 map 函数:

>>> tokenized_wnut = wnut.map(tokenize_and_align_labels, batched=True)

现在使用 DataCollatorWithPadding 创建一个示例批次。在整理数据时,动态填充 句子至批次中最长序列的长度,比将整个数据集填充到最大长度更高效。

>>> from transformers import DataCollatorForTokenClassification

>>> data_collator = DataCollatorForTokenClassification(tokenizer=tokenizer)

评估

在训练过程中包含评估指标通常有助于评估模型的性能。您可以使用 🤗 Evaluate 库快速加载评估方法。对于此任务,请加载 seqeval 框架(请参阅 🤗 Evaluate 的 快速入门 以了解如何加载和计算指标)。Seqeval 实际上会生成多个分数:精确率、召回率、F1 和准确率。

>>> import evaluate

>>> seqeval = evaluate.load("seqeval")

首先获取 NER 标签,然后创建一个函数,将你的真实预测和真实标签传递给 compute 以计算分数:

>>> import numpy as np

>>> labels = [label_list[i] for i in example[f"ner_tags"]]


>>> def compute_metrics(p):
...     predictions, labels = p
...     predictions = np.argmax(predictions, axis=2)

...     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(predictions=true_predictions, references=true_labels)
...     return {
...         "precision": results["overall_precision"],
...         "recall": results["overall_recall"],
...         "f1": results["overall_f1"],
...         "accuracy": results["overall_accuracy"],
...     }

您的 compute_metrics 函数已准备就绪,在配置训练时您将再次用到它。

训练

在开始训练模型之前,请使用 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,
... }

如果您不熟悉使用 Trainer 对模型进行微调,请查看此处的基础教程!

现在可以开始训练模型了!使用 AutoModelForTokenClassification 加载 DistilBERT,并指定预期的标签数量以及标签映射:

>>> from transformers import AutoModelForTokenClassification, TrainingArguments, Trainer

>>> model = AutoModelForTokenClassification.from_pretrained(
...     "distilbert/distilbert-base-uncased", num_labels=13, id2label=id2label, label2id=label2id
... )

此时,只剩下三个步骤:

  1. 在 TrainingArguments 中定义训练超参数。唯一必需的参数是 output_dir,用于指定模型的保存位置。通过设置 push_to_hub=True,你可以将该模型推送到 Hub(上传模型前需要登录 Hugging Face 账户)。在每个 epoch 结束时,Trainer 会评估 seqeval 分数并保存训练检查点。
  2. 将训练参数与模型、数据集、分词器、数据整理器以及 compute_metrics 函数一起传递给 Trainer。
  3. 调用 train() 来微调你的模型。
>>> training_args = TrainingArguments(
...     output_dir="my_awesome_wnut_model",
...     learning_rate=2e-5,
...     per_device_train_batch_size=16,
...     per_device_eval_batch_size=16,
...     num_train_epochs=2,
...     weight_decay=0.01,
...     eval_strategy="epoch",
...     save_strategy="epoch",
...     load_best_model_at_end=True,
...     push_to_hub=True,
... )

>>> trainer = Trainer(
...     model=model,
...     args=training_args,
...     train_dataset=tokenized_wnut["train"],
...     eval_dataset=tokenized_wnut["test"],
...     processing_class=tokenizer,
...     data_collator=data_collator,
...     compute_metrics=compute_metrics,
... )

>>> trainer.train()

训练完成后,使用 push_to_hub() 方法将您的模型共享到 Hub,以便所有人都能使用您的模型:

>>> trainer.push_to_hub()

如需了解如何微调模型以进行 token 分类的更深入示例,请参阅对应的 PyTorch notebook。

推理

很好,现在你已经微调了一个模型,可以使用它进行推理了!

获取一些您想要进行推理的文本:

>>> text = "The Golden State Warriors are an American professional basketball team based in San Francisco."

尝试使用微调后的模型进行推理的最简单方法是将其用于 pipeline()。使用你的模型实例化一个用于 NER 的 pipeline,并将你的文本传入:

>>> from transformers import pipeline

>>> classifier = pipeline("ner", model="stevhliu/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}]

您也可以手动复现 pipeline 的结果:

对文本进行分词并返回 PyTorch 张量:

>>> from transformers import AutoTokenizer

>>> tokenizer = AutoTokenizer.from_pretrained("stevhliu/my_awesome_wnut_model")
>>> inputs = tokenizer(text, return_tensors="pt")

将输入传递给模型并返回 logits:

>>> from transformers import AutoModelForTokenClassification

>>> model = AutoModelForTokenClassification.from_pretrained("stevhliu/my_awesome_wnut_model")
>>> with torch.no_grad():
...     logits = model(**inputs).logits

获取概率最高的类别,并使用模型的 id2label 映射将其转换为文本标签:

>>> predictions = torch.argmax(logits, dim=2)
>>> 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']
Update on GitHub