使用JSON数据集训练GPT-2模型遇Dataset导入错误求助

问题背景 #

使用包含问题、选项、类别及正确答案的JSON数据集训练GPT-2模型时，出现错误：

ImportError: cannot import name 'Dataset' from 'transformers'

用户原代码：

import json
import random
from transformers import GPT2LMHeadModel, Dataset

def convert_json_to_text(json_data):
 

  text = ''
  for question_and_answers in json_data: 
    random.shuffle(question_and_answers['answers'])
 
    text += f"{question_and_answers['category']}: {question_and_answers['question']}\n"
 
    for option in question_and_answers['answers']:
      text += f"- {option}\n"
 
    text += f"Correct Answer: {question_and_answers['correct_answer']}\n\n"

  return text
 
with open("questions.json", "r") as f:
  json_data = json.load(f)
 
text = convert_json_to_text(json_data)
 
train_dataset = Dataset.from_text(text)
 
model = GPT2LMHeadModel.from_pretrained("gpt2")
model.train()
for epoch in range(10):
  for batch in train_dataset:
    loss = model(input_ids=batch['input_ids'], labels=batch['input_ids'])
    loss.backward()
    model.optimizer.step()
    model.optimizer.zero_grad()
 
model.save_pretrained("gpt2_model.pt")

示例数据集：

[
  {
    "question": "Q1. Which operator returns true if the two compared values are not equal?",
    "category": "javascript",
    "answers": [" <>", " ~", " ==!", " !=="],
    "correct_answer": " !=="
  },
  {
    "question": "Q2. How is a forEach statement different from a for statement?",
    "category": "javascript",
    "answers": [
      " Only a for statement uses a callback function.",
      " A for statement is generic, but a forEach statement can be used only with an array.",
      " Only a forEach statement lets you specify your own iterator.",
      " A forEach statement is generic, but a for statement can be used only with an array."
    ],
    "correct_answer": " A for statement is generic, but a forEach statement can be used only with an array."
  }
]

错误原因 #

1. Dataset类不属于transformers库，而是在datasets库中，直接从transformers导入会报错。
2. 原代码缺少必要组件：未使用tokenizer处理文本生成模型所需的input_ids，未初始化优化器，训练循环未正确处理批次数据。

解决步骤 #

1. 安装依赖库 #

首先安装datasets库及其他必要依赖：

pip install datasets transformers torch

2. 修正代码 #

以下是完整的可运行修正代码：

import json
import random
import torch
from datasets import Dataset
from transformers import GPT2LMHeadModel, AutoTokenizer, AdamW, DataCollatorForLanguageModeling
from torch.utils.data import DataLoader

def convert_json_to_examples(json_data):
    examples = []
    for item in json_data:
        random.shuffle(item['answers'])
        # 构建单条训练文本
        text = f"{item['category']}: {item['question']}\n"
        for option in item['answers']:
            text += f"- {option}\n"
        text += f"Correct Answer: {item['correct_answer']}\n\n"
        examples.append({"text": text})
    return examples

# 加载数据集
with open("questions.json", "r") as f:
    json_data = json.load(f)

# 转换为Dataset格式
examples = convert_json_to_examples(json_data)
train_dataset = Dataset.from_list(examples)

# 初始化tokenizer和模型
tokenizer = AutoTokenizer.from_pretrained("gpt2")
tokenizer.pad_token = tokenizer.eos_token  # GPT-2默认没有pad token，用eos token代替

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

# 对数据集进行tokenization处理
tokenized_dataset = train_dataset.map(tokenize_function, batched=True)
tokenized_dataset.set_format("torch", columns=["input_ids", "attention_mask"])

# 数据收集器，自动生成匹配input_ids的labels
data_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer, mlm=False  # GPT是自回归模型，不使用掩码语言建模
)

# 初始化模型和优化器
model = GPT2LMHeadModel.from_pretrained("gpt2")
optimizer = AdamW(model.parameters(), lr=5e-5)

# 训练循环
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
model.train()

epochs = 10
batch_size = 2

# 构建DataLoader处理批次
train_dataloader = DataLoader(tokenized_dataset, batch_size=batch_size, collate_fn=data_collator)

for epoch in range(epochs):
    total_loss = 0
    for batch in train_dataloader:
        batch = {k: v.to(device) for k, v in batch.items()}
        outputs = model(**batch)
        loss = outputs.loss
        total_loss += loss.item()
        
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
    
    print(f"Epoch {epoch+1}, Average Loss: {total_loss / len(train_dataloader)}")

# 保存模型和tokenizer
model.save_pretrained("./trained_gpt2")
tokenizer.save_pretrained("./trained_gpt2")

关键修正点说明 #

• 从datasets库导入Dataset，而非transformers。
• 添加AutoTokenizer处理文本，设置pad token（GPT-2默认无pad token，用eos token替代）。
• 使用DataCollatorForLanguageModeling自动生成labels，匹配input_ids。
• 初始化AdamW优化器，指定学习率。
• 用DataLoader处理批次数据，适配模型训练需求。
• 训练时将数据移至GPU（如果可用），提升训练速度。

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