ModernBERT微调实现文本分类

ModernBERT 是最近发布的传统 BERT 模型（文本嵌入模型）的改进版本。它的表现优于传统 BERT，甚至优于 RoBERTa、ELECTRA 等最佳变体。

在这篇文章中，我们将深入研究针对文本分类任务对 ModernBERT 进行微调。

1、数据集

我们将使用 argilla/synthetic-domain-text-classification 对 ModernBERT 基础模型进行微调。这是一个多类分类数据集，1k 行数据，26 个标签，将文本分类为各个主题。下面是一个快照：

让我们开始吧。

2、安装依赖包

使用 pip 安装所需的软件包：

%pip install "torch==2.5.0" "torchvision==0.20.0"
%pip install "setuptools<71.0.0" scikit-learn

%pip install  --upgrade \
  "datasets==3.1.0" \
  "accelerate==1.2.1" \
  "hf-transfer==0.1.8"

%pip install "git+https://github.com/huggingface/transformers.git@6e0515e99c39444caae39472ee1b2fd76ece32f1" --upgrade

3、加载数据集

让我们加载数据集：

from datasets import load_dataset
from datasets.arrow_dataset import Dataset
from datasets.dataset_dict import DatasetDict, IterableDatasetDict
from datasets.iterable_dataset import IterableDataset

# Dataset id from huggingface.co/dataset
dataset_id = "argilla/synthetic-domain-text-classification"

# Load raw dataset
train_dataset = load_dataset(dataset_id, split='train')

split_dataset = train_dataset.train_test_split(test_size=0.1)
split_dataset['train'][5:7]

4、分词

让我们加载 Tokenizer 并为 ModernBERT 提取数据进行分词：

from transformers import AutoTokenizer

# Model id to load the tokenizer
model_id = "answerdotai/ModernBERT-base"

# Load Tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_id)

# Tokenize helper function
def tokenize(batch):
    return tokenizer(batch['text'], truncation=True,padding=True, return_tensors="pt")

# Tokenize dataset
if "label" in split_dataset["train"].features.keys():
    split_dataset =  split_dataset.rename_column("label", "labels") # to match Trainer
tokenized_dataset = split_dataset.map(tokenize, batched=True, remove_columns=["text"])

以下是代码中发生的情况：

• 设置模型 ID： model_id 指定将为其加载标记器的预训练模型。
• 加载分词器：使用 AutoTokenizer.from_pretrained 加载指定模型的分词器。
• 定义分词器函数：定义辅助函数 tokenize 来处理批量文本数据。它对文本列进行标记，应用截断和填充来标准化序列长度，并转换为 PyTorch 张量。
• 检查并重命名标签列：如果 split_dataset["train"] 有一个标签列，则将其重命名为标签以确保与 Hugging Face 的 Trainer 兼容。
• 对数据集进行分词： map 函数将 tokenize 函数应用于 split_dataset 中的每个批次。它处理文本列并在分词后将其从数据集中删除，从而生成一个可用于训练或评估的数据集。

5、加载ModernBERT预训练模型

在这一部分中，我们将为 id2label 和 label 2id 映射创建几个字典并加载 ModernBERT预训练模型：

from transformers import AutoModelForSequenceClassification

# Model id to load the tokenizer
model_id = "answerdotai/ModernBERT-base"

# Prepare model labels - useful for inference
labels = tokenized_dataset["train"].features["labels"].names
num_labels = len(labels)
label2id, id2label = dict(), dict()
for i, label in enumerate(labels):
    label2id[label] = str(i)
    id2label[str(i)] = label

# Download the model from huggingface.co/models
model = AutoModelForSequenceClassification.from_pretrained(
    model_id, num_labels=num_labels, label2id=label2id, id2label=id2label,
)

这里是代码中发生了什么：

• 设置模型 ID： model_id 指定要加载的预训练模型。
• 准备标签：分类标签是从 tokenized_dataset["train"].features["labels"].names 中提取的。标签总数已确定并存储在 num_labels 中。
• 创建标签映射：创建两个字典 label2id 和 id2label，以将每个标签映射到唯一的数字 ID，反之亦然。此映射是通过迭代标签并为每个标签分配一个 ID 来构建的。
• 加载模型：使用 AutoModelForSequenceClassification.from_pretrained 加载序列分类模型。它使用指定的 model_id 以及 num_labels、 label2id 和 id2label 来配置模型以进行多类分类。这为模型的训练或推理等任务做好了准备。

6、微调ModernBERT

使用如下代码开始微调：

from transformers import Trainer, TrainingArguments

# Define training args
training_args = TrainingArguments(
    output_dir= "ModernBERT-domain-classifier",
    per_device_train_batch_size=4,
    per_device_eval_batch_size=2,
    learning_rate=5e-5,
    num_train_epochs=1,
    bf16=True, # bfloat16 training
    optim="adamw_torch_fused", # improved optimizer
    # logging & evaluation strategies
    logging_strategy="steps",
    eval_strategy="epoch",
    save_strategy="epoch",
    save_total_limit=2,
    load_best_model_at_end=True,
)

# Create a Trainer instance
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset["train"],
    eval_dataset=tokenized_dataset["test"]
)
trainer.train()

使用的一些关键超参数是：

• output_dir：指定训练输出的目录（模型检查点、日志等）将被保存。
• per_device_train_batch_size：训练的批处理大小，设置为 4，这意味着每个设备（例如 GPU）在每个训练步骤中将处理 4 个样本。
• per_device_eval_batch_size：评估的批处理大小，设置为 2，它决定了在评估期间每个设备处理多少个样本。
• learning_rate：初始学习率，设置为 5e-5。这控制每个优化步骤对模型权重的调整程度。
• num_train_epochs：指定整个训练数据集将通过模型的次数，设置为 1。
• bf16：启用使用 bfloat16（Brain Floating Point 16）精度的训练，这可以减少内存使用量并加快训练速度，而不会显着影响准确性。
• optim：指定优化器，设置为 adamw_torch_fused，这是针对性能优化的 AdamW 的改进版本。
• logs_strategy：确定何时生成日志。 steps表示在一定数量的训练步骤后生成日志。
• eval_strategy：定义训练期间评估发生的时间。 epoch表示在每个训练时期结束时进行评估。
• save_strategy：控制何时保存模型检查点。 epoch表示在每个时期结束时保存检查点。
• save_total_limit：将保存的检查点总数限制为 2。一旦超过限制，旧的检查点将被删除。
• load_best_model_at_end：如果设置为 True，则在训练结束时将自动加载最佳模型（基于评估）。

batch_size 和 epochs 具有较低的值以避免内存限制。在对数据进行实际微调时，必须增加此值。

7、利用微调模型进行推理

尝试微调模型：

predictions = trainer.predict(tokenized_dataset["test"])

# Process the prediction results (predictions, label_ids, metrics)
predicted_labels = np.argmax(predictions.predictions, axis=1)
predicted_label = id2label[str(predicted_labels[0])]
example_data = [split_dataset['train'][0]]


print(f"Predicted Label: {predicted_label}")
print(f"Actual Label: {id2label[str(example_data[0]['labels'])]}")

以下是代码片段的作用：

进行预测：

predictions = trainer.predict(tokenized_dataset["test"]) 这一行使用微调模型对测试数据集进行预测。结果包括预测的 logit（原始模型输出）、实际标签和评估指标。

提取预测标签：

np.argmax(predictions.predictions, axis=1) 通过选择每个样本值最高的索引将 logit 转换为预测标签，表示最可能的类别。

映射到标签名称：

id2label[str(predicted_labels[0])] 使用 id2label 字典将第一个预测标签（数字 ID）映射到其对应的标签名称。

选择示例数据：

example_data = [split_dataset['train'][0]] 从训练集中检索第一个数据样本进行比较。

打印预测和实际标签：

使用 print 打印第一个测试样本的预测标签和所选训练样本的实际标签。

这样就大功告成了。尝试使用您的自定义数据并微调 ModernBERT。

原文链接：ModernBERT for Text Classification

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