如何修改Hugging Face预训练模型分类标签？适配自定义多分类任务

核心思路 #

原模型cardiffnlp/twitter-roberta-base-sentiment-latest针对3类情感（Positive/Negative/Neutral）训练，要适配4类自定义标签（satisfied/positive/frustrated/confused），最优方案是复用预训练RoBERTa主体，替换并微调新的4分类头——既保留预训练的语义特征，又快速适配新分类任务，数据集较小时还可冻结主体层避免过拟合。

具体步骤 #

1. 准备标注好的自定义数据集：每条数据包含文本和对应标签（映射为0-3的整数）
2. 加载预训练模型与分词器，修改分类层输出维度为4
3. 配置训练参数、损失函数与优化器，完成模型微调
4. 基于微调后的模型实现推理逻辑

完整代码实现 #

1. 依赖安装 #

pip install transformers datasets torch numpy scikit-learn

2. 数据集准备示例 #

假设自定义数据集为CSV格式，包含text（文本）和label（标签列，0=satisfied、1=positive、2=frustrated、3=confused）：

from datasets import load_dataset

# 加载自定义数据集（替换为你的数据集路径）
dataset = load_dataset('csv', data_files='custom_sentiment_data.csv')

# 定义标签映射
label2id = {"satisfied":0, "positive":1, "frustrated":2, "confused":3}
id2label = {v:k for k,v in label2id.items()}

3. 模型加载与微调代码 #

import torch
import numpy as np
from transformers import (
    AutoTokenizer, AutoModelForSequenceClassification,
    TrainingArguments, Trainer, DataCollatorWithPadding
)
from sklearn.metrics import accuracy_score, f1_score

# 加载预训练模型与分词器
MODEL_NAME = "cardiffnlp/twitter-roberta-base-sentiment-latest"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

# 替换分类头为4分类
model = AutoModelForSequenceClassification.from_pretrained(
    MODEL_NAME,
    num_labels=4,
    id2label=id2label,
    label2id=label2id
)

# 数据集预处理函数
def preprocess_function(examples):
    return tokenizer(examples["text"], truncation=True, max_length=128)

# 批量预处理数据集
tokenized_dataset = dataset.map(preprocess_function, batched=True)
data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

# 定义评估指标
def compute_metrics(eval_pred):
    logits, labels = eval_pred
    predictions = np.argmax(logits, axis=-1)
    return {
        "accuracy": accuracy_score(labels, predictions),
        "f1_weighted": f1_score(labels, predictions, average="weighted")
    }

# 训练参数配置
training_args = TrainingArguments(
    output_dir="./custom_sentiment_model",
    learning_rate=2e-5,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=16,
    num_train_epochs=3,
    weight_decay=0.01,
    evaluation_strategy="epoch",
    save_strategy="epoch",
    load_best_model_at_end=True,
    logging_dir="./logs"
)

# 初始化训练器
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset["train"],
    eval_dataset=tokenized_dataset["test"], # 假设数据集拆分有test集
    tokenizer=tokenizer,
    data_collator=data_collator,
    compute_metrics=compute_metrics,
)

# 启动训练
trainer.train()

# 保存最终模型
trainer.save_model("./custom_sentiment_final")

4. 自定义分类推理代码 #

import torch
import numpy as np
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from scipy.special import softmax

def classify_custom_sentiment(text):
    # 加载微调后的模型（替换为你的模型保存路径）
    MODEL_PATH = "./custom_sentiment_final"
    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
    model = AutoModelForSequenceClassification.from_pretrained(MODEL_PATH)
    config = model.config

    # 文本编码与推理
    encoded_input = tokenizer(text, return_tensors='pt', truncation=True, max_length=128)
    with torch.no_grad():
        output = model(**encoded_input)
    
    scores = output[0][0].detach().numpy()
    scores = softmax(scores)
    ranking = np.argsort(scores)[::-1]

    results = {}
    for i in range(scores.shape[0]):
        label = config.id2label[ranking[i]]
        score = np.round(float(scores[ranking[i]]), 4)
        results[label] = score

    return results

# 测试示例
print(classify_custom_sentiment("这个产品的功能完全符合我的预期，很满意"))
print(classify_custom_sentiment("完全搞不懂这个操作步骤，越用越懵"))

关键说明 #

• 若数据集规模较小，可添加model.roberta.requires_grad_(False)冻结预训练主体，仅训练分类头，避免过拟合
• 可根据数据情况调整训练参数（如学习率、批次大小、训练轮数）
• 可在预处理阶段加入文本清洗逻辑（如去除特殊字符、统一大小写），进一步提升模型效果

内容的提问来源于stack exchange，提问作者Gangadhar Neelam