2025年NLP神器 - Hugging Face - 预训练模型微调demo

第一个完整的使用Hugging Face Transformers对预训练模型进行微调的demo is coming！

整体步骤为：

加载数据集load dataset

from datasets import load_dataset
raw_datasets = load_dataset("glue", "mrpc")

对数据集做分词tokenize

from transformers import AutoModelForSequenceClassification
checkpoint = 'distilbert-base-uncased-finetuned-sst-2-english'
tokenizer = AutoTokenizer.from_pretrained(checkpoint)
 

def tokenize_function(examples):
    return tokenizer(examples["sentence1"], examples["sentence2"],truncation=True)
 
tokenized_datasets = raw_datasets.map(tokenize_function, batched=True)

组装DataCollator

from transformers import DataCollatorWithPadding
 
data_collator = DataCollatorWithPadding(tokenizer, padding=True)

构建模型model

model = AutoModelForSequenceClassification.from_pretrained(checkpoint, num_labels=2)

定义模型参数TrainingArgs

from transformers import TrainingArguments
training_args = TrainingArguments('test-trainer')

创建训练对象Trainer

from transformers import Trainer
trainer = Trainer(
    model,
    training_args,
    train_dataset=tokenized_datasets['train'],
    eval_dataset=tokenized_datasets['validation'],
    data_collator=data_collator,
    tokenizer=tokenizer,
)

训练 ---  一键开启，简便至极

trainer.train()

预测

predictions = trainer.predict(tokenized_datasets['test'])
print(predictions.predictions.shape, predictions.label_ids.shape)

import numpy as np
 
preds = np.argmax(predictions.predictions, axis=-1)

评估 --- 最新版本的评估模块，对得起标题2025

import evaluate
 
accuracy = evaluate.load('accuracy')
print(accuracy.compute(references=predictions.label_ids, predictions=preds))
clf_metrics = evaluate.combine(["accuracy", "f1", "precision", "recall"])
all_res = clf_metrics.compute(predictions=preds, references=predictions.label_ids)
print(all_res)

上完整代码

# 使用load dataset方法获取公开数据集。这里使用的是GLUE下的MRPC数据集，包含两短文本，和这两个文本的关系是否相关（标签数量为2）
from datasets import load_dataset
raw_datasets = load_dataset("glue", "mrpc")

# 获取分词器tokenizer
from transformers import AutoModelForSequenceClassification
checkpoint = 'distilbert-base-uncased-finetuned-sst-2-english'
tokenizer = AutoTokenizer.from_pretrained(checkpoint)

# 对数据集进行处理时，Hugging Face官方推荐使用map函数，如果不习惯也可以采用pandas的处理方式
# 此时需要预先定一个转换函数给map使用
def tokenize_function(examples):
    return tokenizer(examples["sentence1"], examples["sentence2"],truncation=True)

tokenized_datasets = raw_datasets.map(tokenize_function, batched=True)

# DataCollator其实就是PyTorch里的Dataloader，按批次组装数据集，然后输入给模型
from transformers import DataCollatorWithPadding # regard it as dataloader in pytorch

data_collator = DataCollatorWithPadding(tokenizer, padding=True)

# 获取到预训练模型后，这里指定输出类别为2，相当于自定义给该预训练模型上加了一个输出头，做分类输出
# 这段代码运行后会看到提示，这个输出头部分的权重无法直接从预训练模型中加载，需要后面我们自己来进行训练
model = AutoModelForSequenceClassification.from_pretrained(checkpoint, num_labels=2)

# 构建训练参数，参数非常多，这里选择一个简单的通用式
from transformers import TrainingArguments
training_args = TrainingArguments('test-trainer')

# 构建训练对象
from transformers import Trainer
trainer = Trainer(
    model,
    training_args,
    train_dataset=tokenized_datasets['train'],
    eval_dataset=tokenized_datasets['validation'],
    data_collator=data_collator,
    tokenizer=tokenizer,
)

# 开始训练
trainer.train()

# 预测
predictions = trainer.predict(tokenized_datasets['test'])
print(predictions.predictions.shape, predictions.label_ids.shape)
# (1725, 2) (1725,)

import numpy as np

# 解析预测结果
preds = np.argmax(predictions.predictions, axis=-1)

# 评估一下
import evaluate

accuracy = evaluate.load('accuracy')
print(accuracy.compute(references=predictions.label_ids, predictions=preds))
clf_metrics = evaluate.combine(["accuracy", "f1", "precision", "recall"])
all_res = clf_metrics.compute(predictions=preds, references=predictions.label_ids)
print(all_res)