请你再加入训练的代码和网格搜索的代码

好的，以下是加入训练的代码：

from transformers import GPT2LMHeadModel, GPT2Tokenizer, Trainer, TrainingArguments
import torch

tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
model = GPT2LMHeadModel.from_pretrained('gpt2')

# 加载数据集
train_dataset = torch.load('train_dataset.pth')
valid_dataset = torch.load('valid_dataset.pth')

# 训练参数
training_args = TrainingArguments(
    output_dir='./results',
    num_train_epochs=2,
    per_device_train_batch_size=8,
    per_device_eval_batch_size=8,
    warmup_steps=500,
    weight_decay=0.01,
    logging_dir='./logs',
    logging_steps=500,
    evaluation_strategy='steps',
    eval_steps=1000,
    save_total_limit=3,
    save_steps=1000,
    learning_rate=2e-5,
)

# 训练模型
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=valid_dataset,
    data_collator=lambda data: {'input_ids': torch.stack([f[0] for f in data]),
                                'attention_mask': torch.stack([f[1] for f in data]),
                                'labels': torch.stack([f[0] for f in data])}
)

trainer.train()

以下是网格搜索的代码：

from transformers import GPT2LMHeadModel, GPT2Tokenizer, Trainer, TrainingArguments
from sklearn.model_selection import ParameterGrid
import torch

tokenizer = GPT2Tokenizer.from_pretrained('gpt2')

# 加载数据集
train_dataset = torch.load('train_dataset.pth')
valid_dataset = torch.load('valid_dataset.pth')

# 训练参数
params = {'learning_rate': [2e-5, 3e-5],
          'num_train_epochs': [2, 3],
          'per_device_train_batch_size': [8, 16],
          'warmup_steps': [500, 1000]}

# 网格搜索
for param in ParameterGrid(params):
    print(param)
    training_args = TrainingArguments(
        output_dir='./results',
        num_train_epochs=param['num_train_epochs'],
        per_device_train_batch_size=param['per_device_train_batch_size'],
        per_device_eval_batch_size=param['per_device_train_batch_size'],
        warmup_steps=param['warmup_steps'],
        weight_decay=0.01,
        logging_dir='./logs',
        logging_steps=500,
        evaluation_strategy='steps',
        eval_steps=1000,
        save_total_limit=3,
        save_steps=1000,
        learning_rate=param['learning_rate'],
    )

    model = GPT2LMHeadModel.from_pretrained('gpt2')
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=valid_dataset,
        data_collator=lambda data: {'input_ids': torch.stack([f[0] for f in data]),
                                    'attention_mask': torch.stack([f[1] for f in data]),
                                    'labels': torch.stack([f[0] for f in data])}
    )

    trainer.train()