文本分类模型教程

文本分类是一种常见的自然语言处理任务,将标签或类别分配给文本。一些大公司在实际应用中运行文本分类,用于各种实际应用。最流行的文本分类形式之一是情感分析,将标签(如 🙂 positive,🙁 negative 或 😐 neutral)分配给一系列文本。

本指南将展示如何:

  1. IMDb 数据集上使用 DistilBERT 进行微调,确定电影评论是积极还是消极。
  2. 使用微调后的模型进行推理。

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

pip install transformers datasets evaluate

我们建议您登录到您的 Hugging Face 帐户,这样您就可以上传和与社区共享您的模型。提示时,输入您的令牌进行登录:

>>> from huggingface_hub import notebook_login

>>> notebook_login()

加载 IMDb 数据集

首先从 🤗 数据集库中加载 IMDb 数据集:

>>> from datasets import load_dataset

>>> imdb = load_dataset("imdb")

然后查看一个示例:

>>> imdb["test"][0]
{
    "label": 0,
    "text": "I love sci-fi and am willing to put up with a lot. Sci-fi movies/TV are usually underfunded, under-appreciated and misunderstood. I tried to like this, I really did, but it is to good TV sci-fi as Babylon 5 is to Star Trek (the original). Silly prosthetics, cheap cardboard sets, stilted dialogues, CG that doesn't match the background, and painfully one-dimensional characters cannot be overcome with a 'sci-fi' setting. (I'm sure there are those of you out there who think Babylon 5 is good sci-fi TV. It's not. It's clichéd and uninspiring.) While US viewers might like emotion and character development, sci-fi is a genre that does not take itself seriously (cf. Star Trek). It may treat important issues, yet not as a serious philosophy. It's really difficult to care about the characters here as they are not simply foolish, just missing a spark of life. Their actions and reactions are wooden and predictable, often painful to watch. The makers of Earth KNOW it's rubbish as they have to always say \"Gene Roddenberry's Earth...\" otherwise people would not continue watching. Roddenberry's ashes must be turning in their orbit as this dull, cheap, poorly edited (watching it without advert breaks really brings this home) trudging Trabant of a show lumbers into space. Spoiler. So, kill off a main character. And then bring him back as another actor. Jeeez! Dallas all over again.",
}

这个数据集有两个字段:

  • text:电影评论文本。
  • label:0 表示负面评论,1 表示正面评论。

预处理

下一步是加载一个 DistilBERT 的 tokenizer 来预处理 text 字段:

>>> from transformers import AutoTokenizer

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

创建一个预处理函数将 text 进行分词,并截断超过 DistilBERT 最大输入长度的序列:

>>> def preprocess_function(examples):
...     return tokenizer(examples["text"], truncation=True)

使用 🤗 数据集 map 函数将预处理函数应用于整个数据集,通过将 batched=True 设置来一次处理数据集的多个元素,可以加速 map 的处理速度:

tokenized_imdb = imdb.map(preprocess_function, batched=True)

现在使用 DataCollatorWithPadding 创建一批样例。在整个数据集上进行填充,以整个批次中最长句子的长度为最大长度,比每个样本单独填充要高效得多。

Pytorch
Hide Pytorch content
>>> from transformers import DataCollatorWithPadding

>>> data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
TensorFlow
Hide TensorFlow content
>>> from transformers import DataCollatorWithPadding

>>> data_collator = DataCollatorWithPadding(tokenizer=tokenizer, return_tensors="tf")

评估

在训练过程中包含一个评估指标通常有助于评估模型的表现。您可以使用 🤗 Evaluate 库快速加载一个评估方法。对于此任务,加载 accuracy 指标(请参阅 🤗 Evaluate 的 快速导览 以了解有关如何加载和计算指标的更多信息):

>>> import evaluate

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

然后创建一个函数,将预测值和标签传递给 compute 来计算准确率:

>>> import numpy as np


>>> def compute_metrics(eval_pred):
...     predictions, labels = eval_pred
...     predictions = np.argmax(predictions, axis=1)
...     return accuracy.compute(predictions=predictions, references=labels)

您的 compute_metrics 函数已经准备好了,在设置训练时会再次使用。

训练

在开始模型训练之前,使用 id2labellabel2id 创建一个期望的标签到其 id 的映射集合:

>>> id2label = {0: "NEGATIVE", 1: "POSITIVE"}
>>> label2id = {"NEGATIVE": 0, "POSITIVE": 1}
Pytorch
Hide Pytorch content

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

你现在可以开始训练你的模型了!使用 `AutoModelForSequenceClassification` 加载 DistilBERT,同时指定期望的标签数量和标签映射关系:
>>> from transformers import AutoModelForSequenceClassification, TrainingArguments, Trainer

>>> model = AutoModelForSequenceClassification.from_pretrained(
...     "distilbert-base-uncased", num_labels=2, id2label=id2label, label2id=label2id
... )

在这一步中,还需要完成以下三个步骤:

  1. TrainingArguments中定义你的训练超参数。唯一需要的参数是output_dir,用于指定保存模型的位置。你可以通过设置push_to_hub=True将该模型上传到 Hub(上传模型需要登录 Hugging Face)。在每个 epoch 结束时,Trainer将评估准确性并保存训练检查点。
  2. 将训练参数和模型、数据集、分词器、数据收集器和compute_metrics函数传递给Trainer
  3. 调用train()对模型进行微调。
>>> training_args = TrainingArguments(
...     output_dir="my_awesome_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,
...     evaluation_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_imdb["train"],
...     eval_dataset=tokenized_imdb["test"],
...     tokenizer=tokenizer,
...     data_collator=data_collator,
...     compute_metrics=compute_metrics,
... )

>>> trainer.train()

当你将tokenizer传递给Trainer时,默认情况下Trainer会应用动态填充。在这种情况下,你不需要显式指定数据收集器(data collator)。

训练完成后,使用push_to_hub()方法将模型分享到 Hub,这样其他人就可以使用你的模型了:

>>> trainer.push_to_hub()
TensorFlow
Hide TensorFlow content

如果你对使用 Keras 进行微调的模型不熟悉,可以参考 这里 的基础教程!

在 TensorFlow 中微调模型时,首先要设置优化器函数、学习率调度和一些训练超参数:

>>> from transformers import create_optimizer
>>> import tensorflow as tf

>>> batch_size = 16
>>> num_epochs = 5
>>> batches_per_epoch = len(tokenized_imdb["train"]) // batch_size
>>> total_train_steps = int(batches_per_epoch * num_epochs)
>>> optimizer, schedule = create_optimizer(init_lr=2e-5, num_warmup_steps=0, num_train_steps=total_train_steps)

然后你可以使用TFAutoModelForSequenceClassification加载 DistilBERT,同时指定期望的标签数量和标签映射关系:

>>> from transformers import TFAutoModelForSequenceClassification

>>> model = TFAutoModelForSequenceClassification.from_pretrained(
...     "distilbert-base-uncased", num_labels=2, id2label=id2label, label2id=label2id
... )

通过prepare_tf_dataset()将数据集转换为tf.data.Dataset格式:

>>> tf_train_set = model.prepare_tf_dataset(
...     tokenized_imdb["train"],
...     shuffle=True,
...     batch_size=16,
...     collate_fn=data_collator,
... )

>>> tf_validation_set = model.prepare_tf_dataset(
...     tokenized_imdb["test"],
...     shuffle=False,
...     batch_size=16,
...     collate_fn=data_collator,
... )

使用 compile 配置模型进行训练。注意,Transformers 模型都具有默认的与任务相关的损失函数,因此除非你想要自定义损失函数,否则不需要指定:

>>> import tensorflow as tf

>>> model.compile(optimizer=optimizer)  # 没有指定损失参数!

在开始训练之前,还需要计算预测的准确性,并提供将模型上传到 Hub 的方式。这两个任务都可以通过使用 Keras 回调 来完成。

将你的compute_metrics函数传递给KerasMetricCallback

>>> from transformers.keras_callbacks import KerasMetricCallback

>>> metric_callback = KerasMetricCallback(metric_fn=compute_metrics, eval_dataset=tf_validation_set)

PushToHubCallback中指定模型和分词器的上传位置:

>>> from transformers.keras_callbacks import PushToHubCallback

>>> push_to_hub_callback = PushToHubCallback(
...     output_dir="my_awesome_model",
...     tokenizer=tokenizer,
... )

然后将所有回调函数封装在一起:

>>> callbacks = [metric_callback, push_to_hub_callback]

最后,你可以开始训练模型了!通过 fit 方法传递你的训练和验证数据集、epoch 数以及回调函数来进行模型微调:

>>> model.fit(x=tf_train_set, validation_data=tf_validation_set, epochs=3, callbacks=callbacks)

训练完成后,你的模型将自动上传到 Hub,这样每个人都可以使用它!

若要了解有关如何为文本分类微调模型的更详细示例,请参考相应的 PyTorch notebookTensorFlow notebook

推理

非常好,现在你已经对模型进行了微调,可以开始使用它进行推理了!

选择一些你想要进行推理的文本:

>>> text = "This was a masterpiece. Not completely faithful to the books, but enthralling from beginning to end. Might be my favorite of the three."

使用最简单的方式对模型进行推理是使用pipeline()。用你的模型实例化一个情感分析的pipeline,并将文本传递给它:

>>> from transformers import pipeline

>>> classifier = pipeline("sentiment-analysis", model="stevhliu/my_awesome_model")
>>> classifier(text)
[{'label': 'POSITIVE', 'score': 0.9994940757751465}]

您也可以自己手动实现 pipeline 并得到一致的结果:

Pytorch
Hide Pytorch content

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

>>> from transformers import AutoTokenizer

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

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

>>> from transformers import AutoModelForSequenceClassification

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

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

>>> predicted_class_id = logits.argmax().item()
>>> model.config.id2label[predicted_class_id]
'POSITIVE'
TensorFlow
Hide TensorFlow content

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

>>> from transformers import AutoTokenizer

>>> tokenizer = AutoTokenizer.from_pretrained("stevhliu/my_awesome_model")
>>> inputs = tokenizer(text, return_tensors="tf")

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

>>> from transformers import TFAutoModelForSequenceClassification

>>> model = TFAutoModelForSequenceClassification.from_pretrained("stevhliu/my_awesome_model")
>>> logits = model(**inputs).logits

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

>>> predicted_class_id = int(tf.math.argmax(logits, axis=-1)[0])
>>> model.config.id2label[predicted_class_id]
'POSITIVE'