一起学Hugging Face Transformers（5）- 使用Transformers 库进行文本分类

前言

文本分类（Text Classification）是自然语言处理（NLP）中的一个基本任务，它涉及将文本分配到预定义的类别中。Hugging Face 的 Transformers 库为文本分类提供了强大而简便的工具。本文将详细介绍如何使用 Transformers 库进行文本分类，包括数据预处理、模型训练和评估。

一、环境准备与代码编写

1. 环境准备

首先，确保安装了必要的库：

pip install transformers datasets torch

2. 导入库和加载数据

在本示例中，我们将使用 Hugging Face 的 datasets 库来加载示例数据集。我们使用 IMDb 影评数据集进行二分类任务（正面评价和负面评价）。

from datasets import load_dataset

# 加载 IMDb 数据集
dataset = load_dataset('imdb')

3. 数据预处理

Tokenization

Tokenization 是将文本转换为模型可以处理的数字序列。在 Hugging Face 的 Transformers 库中，AutoTokenizer 类可以自动选择并加载合适的分词器（Tokenizer）。

from transformers import AutoTokenizer

# 选择一个预训练模型的分词器，例如 BERT
tokenizer = AutoTokenizer.from_pretrained('bert-base-uncased')

# 定义一个函数进行数据预处理
def preprocess_function(examples):
    return tokenizer(examples['text'], padding='max_length', truncation=True, max_length=128)

# 对数据集进行 Tokenization
encoded_dataset = dataset.map(preprocess_function, batched=True)

在这个示例中，我们使用 BERT 的分词器，并将每个文本转换为最大长度为128的数字序列。

4. 创建数据集和 DataLoader

接下来，我们需要将数据转换为适合 PyTorch 模型训练的数据格式。

import torch
from torch.utils.data import DataLoader

# 转换数据格式
encoded_dataset.set_format(type='torch', columns=['input_ids', 'attention_mask', 'label'])

# 创建数据加载器
train_dataloader = DataLoader(encoded_dataset['train'], batch_size=8, shuffle=True)
eval_dataloader = DataLoader(encoded_dataset['test'], batch_size=8)

5. 加载预训练模型

我们使用一个预训练的 Transformer 模型，并为文本分类任务进行微调（fine-tuning）。这里我们使用 BERT 模型。

from transformers import AutoModelForSequenceClassification, TrainingArguments, Trainer

# 加载预训练模型
model = AutoModelForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=2)

6. 定义训练参数和训练模型

我们使用 Trainer 类来简化训练过程。首先，我们需要定义训练参数。

training_args = TrainingArguments(
    output_dir='./results',
    evaluation_strategy='epoch',
    learning_rate=2e-5,
    per_device_train_batch_size=8,
    per_device_eval_batch_size=8,
    num_train_epochs=3,
    weight_decay=0.01
)

# 定义 Trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=encoded_dataset['train'],
    eval_dataset=encoded_dataset['test']
)

# 开始训练
trainer.train()

7. 评估模型

训练完成后，我们可以使用 Trainer 提供的 evaluate 方法对模型进行评估。

# 评估模型
eval_results = trainer.evaluate()

print(f"Evaluation results: {eval_results}")

8. 模型推理

我们可以使用训练好的模型对新文本进行分类。

# 定义推理函数
def predict(text):
    inputs = tokenizer(text, return_tensors='pt', padding=True, truncation=True, max_length=128)
    outputs = model(**inputs)
    predictions = torch.argmax(outputs.logits, dim=-1)
    return predictions

# 示例推理
texts = ["I love this movie!", "This movie is terrible."]
predictions = predict(texts)

print(predictions)

二、运行与跟踪

在这个示例中，我们通过使用 Hugging Face 的 Transformers 库对 IMDb 影评数据集进行了文本分类任务。在实际运行过程中，你会得到如下类型的结果：

1. 数据预处理和加载数据

在运行以下代码时：

from datasets import load_dataset

# 2.加载 IMDb 数据集
dataset = load_dataset('imdb')

你会看到加载 IMDb 数据集的进度信息，如：

Downloading: 100%|███████████████████████████████████████████████| 25.0k/25.0k [00:00<00:00, 25.4MB/s]
...

2.Tokenization

在对数据集进行 Tokenization 时：

encoded_dataset = dataset.map(preprocess_function, batched=True)

你会看到数据处理的进度，如：

  0%|          | 0/25 [00:00<?, ?ba/s]
100%|██████████| 25/25 [00:05<00:00,  4.52ba/s]

3.创建数据集和 DataLoader

在转换数据格式并创建 DataLoader 时：

# 转换数据格式
encoded_dataset.set_format(type='torch', columns=['input_ids', 'attention_mask', 'label'])
# 创建数据加载器
train_dataloader = DataLoader(encoded_dataset['train'], batch_size=8, shuffle=True)
eval_dataloader = DataLoader(encoded_dataset['test'], batch_size=8)

这些步骤通常不会输出太多信息，但会确保数据格式正确。

4.加载预训练模型

在加载预训练模型时：

model = AutoModelForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=2)

你会看到模型下载和加载的进度信息，如：

Downloading: 100%|███████████████████████████████████████████████| 420M/420M [00:12<00:00, 33.9MB/s]
...

5.模型训练

在训练模型时：

trainer.train()

你会看到训练过程中的日志输出，包括每个 epoch 的训练损失（loss）和评估指标（evaluation metrics）：

***** Running training *****
  Num examples = 25000
  Num Epochs = 3
  Instantaneous batch size per device = 8
  Total train batch size (w. parallel, distributed & accumulation) = 8
  Gradient Accumulation steps = 1
  Total optimization steps = 9375
...

6.评估模型

在评估模型时：

eval_results = trainer.evaluate()
print(f"Evaluation results: {eval_results}")

你会得到评估结果，如：

***** Running Evaluation *****
  Num examples = 25000
  Batch size = 8
...
Evaluation results: {'eval_loss': 0.405, 'eval_accuracy': 0.855}

这里的结果示例表示评估损失为0.405，准确率为85.5%。

7.模型推理

在对新文本进行推理时：

texts = ["I love this movie!", "This movie is terrible."]
predictions = predict(texts)
print(predictions)

你会得到模型的预测结果，如：

tensor([1, 0])

这里的结果示例表示第一条评论（“I love this movie!”）被分类为类别1（可能是正面评价），第二条评论（“This movie is terrible.”）被分类为类别0（可能是负面评价）。

总结

通过上述步骤，我们可以使用 Hugging Face 的 Transformers 库轻松地进行文本分类。该库提供了丰富的预训练模型和简便的 API，使得从数据预处理到模型训练和评估的全过程变得非常简洁高效。希望本文能帮助你更好地理解和应用 Transformers 库进行文本分类任务。