from transformers import GPT2LMHeadModel, GPT2Tokenizer, Trainer, TrainingArguments
from datasets import load_dataset, Dataset
import json

# 加载数据集
def load_filtered_data(file_path):
    with open(file_path, 'r', encoding='utf-8') as f:
        data = json.load(f)
    texts = []
    for entry in data:
        texts.extend([entry[key] for key in entry if entry[key]])
    return Dataset.from_dict({"text": texts})

# 定义模型和标记器路径
model_path = "需要微调的模型路径"
tokenizer = GPT2Tokenizer.from_pretrained(model_path)
model = GPT2LMHeadModel.from_pretrained(model_path)

# 添加填充标记
tokenizer.pad_token = tokenizer.eos_token

# 加载数据集
dataset = load_filtered_data("D:/data/数据集.json")

def tokenize_function(examples):
    # 对输入文本进行标记化处理
    return tokenizer(examples['text'], truncation=True, padding='max_length', max_length=512)

tokenized_dataset = dataset.map(tokenize_function, batched=True)

# 设置labels
def set_labels(examples):
    examples['labels'] = examples['input_ids'].copy()
    return examples

tokenized_dataset = tokenized_dataset.map(set_labels, batched=True)

# 定义训练参数
training_args = TrainingArguments(
    output_dir="D:/pycode/results",
    overwrite_output_dir=True,
    num_train_epochs=3,
    per_device_train_batch_size=2,
    save_steps=10,
    save_total_limit=2,
    logging_dir="D:/pycode/logs",
    logging_steps=10,
)

# 创建Trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset,
)

# 开始训练
trainer.train()