Transformers 文档

文本分类

Hugging Face's logo
加入Hugging Face社区

并获得增强文档体验

开始使用

文本分类

文本分类是常见的NLP任务,它将标签或类别分配给文本。一些大型公司在生产环境中运行文本分类,用于各种实际应用。最流行的文本分类形式之一是情感分析,它将标签(例如🙂正面、🙁负面或😐中性)分配给一段文本。

本指南将向您展示如何

  1. DistilBERT 上微调 IMDb 数据集以确定电影评论是正面还是负面。
  2. 将您微调的模型用于推理。

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

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

pip install transformers datasets evaluate accelerate

我们鼓励您登录您的Hugging Face帐户,以便您可以上传并与社区共享您的模型。当系统提示时,输入您的令牌以登录

>>> from huggingface_hub import notebook_login

>>> notebook_login()

加载IMDb数据集

首先从 🤗 Datasets 库加载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分词器以预处理text字段

>>> from transformers import AutoTokenizer

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

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

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

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

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

现在使用 DataCollatorWithPadding 创建一批示例。在整理时,动态填充句子以匹配批次中的最长长度,而不是将整个数据集填充到最大长度,这样效率更高。

Pytorch
隐藏Pytorch内容
>>> from transformers import DataCollatorWithPadding

>>> data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
TensorFlow
隐藏TensorFlow内容
>>> from transformers import DataCollatorWithPadding

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

评估

在训练期间包含指标通常有助于评估模型的性能。您可以使用 🤗 Evaluate 库快速加载评估方法。对于此任务,加载 精度 指标(参见 🤗 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
隐藏Pytorch内容

如果您不熟悉使用 Trainer 微调模型,请查看基本教程 此处

现在您可以开始训练您的模型了!使用 AutoModelForSequenceClassification 加载DistilBERT,以及预期的标签数量和标签映射

>>> from transformers import AutoModelForSequenceClassification, TrainingArguments, Trainer

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

此时,只剩下三个步骤

  1. TrainingArguments 中定义您的训练超参数。唯一必需的参数是output_dir,它指定保存模型的位置。您将通过设置push_to_hub=True将此模型推送到Hub(您需要登录Hugging Face才能上传您的模型)。在每个纪元结束时,Trainer 将评估精度并保存训练检查点。
  2. 将训练参数传递给 Trainer,以及模型、数据集、分词器、数据整理器和compute_metrics函数。
  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,
...     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_imdb["train"],
...     eval_dataset=tokenized_imdb["test"],
...     tokenizer=tokenizer,
...     data_collator=data_collator,
...     compute_metrics=compute_metrics,
... )

>>> trainer.train()

当您将tokenizer传递给 Trainer 时,它默认情况下会应用动态填充。在这种情况下,您不需要显式指定数据整理器。

训练完成后,使用 push_to_hub() 方法将您的模型共享到Hub,以便每个人都可以使用您的模型

>>> trainer.push_to_hub()
TensorFlow
隐藏TensorFlow内容

如果您不熟悉使用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/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)  # No loss argument!

在开始训练之前,还需要设置两件事:从预测中计算精度,以及提供一种将模型推送到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 微调模型

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

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

有关如何微调文本分类模型的更深入示例,请查看相应的 PyTorch笔记本TensorFlow笔记本

推理

很好,现在您已经微调了模型,可以将其用于推理了!

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

>>> 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
隐藏Pytorch内容

对文本进行分词并返回 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
隐藏TensorFlow内容

对文本进行分词并返回 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'
< > 在 GitHub 上更新