如何基于多个无关数据集训练Hugging Face模型？

你提到的interleave_datasets确实是解决这类多数据集混合训练的合适工具，不过需要先统一两个数据集的格式，确保它们能被模型正确处理。以下是修改后的完整方案：

核心步骤 #

• 对Short Jokes数据集做tokenize处理，同时补充序列分类任务必需的label字段
• 使用interleave_datasets按指定比例混合两个数据集的训练集和验证集
• 将混合后的数据集传入Trainer完成训练

修改后的完整代码 #

from datasets import load_dataset, interleave_datasets

# 加载两个目标数据集
yelp_dataset = load_dataset("yelp_review_full")
jokes_dataset = load_dataset("short-jokes")

from transformers import AutoTokenizer, Trainer

tokenizer = AutoTokenizer.from_pretrained("google-bert/bert-base-cased")

def tokenize_function(examples):
    return tokenizer(examples["text"], padding="max_length", truncation=True)

# 处理Yelp评论数据集
tokenized_yelp = yelp_dataset.map(tokenize_function, batched=True)
small_train_yelp = tokenized_yelp["train"].shuffle(seed=42).select(range(1000))
small_eval_yelp = tokenized_yelp["test"].shuffle(seed=42).select(range(1000))

# 处理Short Jokes数据集：添加label字段（这里统一设为0，需保证在模型num_labels的范围内）
def add_dummy_label(examples):
    examples["label"] = [0] * len(examples["text"])
    return examples

# 先添加标签字段，再执行tokenize
tokenized_jokes = jokes_dataset.map(add_dummy_label, batched=True).map(tokenize_function, batched=True)
small_train_jokes = tokenized_jokes["train"].shuffle(seed=42).select(range(1000))
# 注意：short-jokes原生没有test集，这里用训练集的子集替代
small_eval_jokes = tokenized_jokes["train"].shuffle(seed=42).select(range(1000))

# 混合训练集与验证集，probabilities参数控制两个数据集的采样比例
mixed_train_dataset = interleave_datasets([small_train_yelp, small_train_jokes], probabilities=[0.5, 0.5], seed=42)
mixed_eval_dataset = interleave_datasets([small_eval_yelp, small_eval_jokes], probabilities=[0.5, 0.5], seed=42)

from transformers import AutoModelForSequenceClassification

model = AutoModelForSequenceClassification.from_pretrained("google-bert/bert-base-cased", num_labels=5)

from transformers import TrainingArguments

training_args = TrainingArguments(output_dir="test_trainer")

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=mixed_train_dataset,
    eval_dataset=mixed_eval_dataset,
)

trainer.train()

关键注意事项 #

• 字段一致性：两个数据集必须包含模型所需的全部输入字段（此处为input_ids、attention_mask、label），因此给无标签的Short Jokes添加了虚拟标签
• 采样比例调整：通过interleave_datasets的probabilities参数可以灵活控制两个数据集的训练占比，比如想侧重Yelp数据可设为[0.7, 0.3]
• 数据集划分：short-jokes没有官方划分的test集，可自行拆分训练集或用现有子集替代
• 任务适配：由于你使用的是序列分类模型，所以将两个数据集统一为分类任务；若需适配不同任务，需额外设置多任务损失逻辑

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