import torch
from transformers import LlamaTokenizer
from model import *
from config import *
from data import *
from torch.nn import CrossEntropyLoss
from torch.cuda.amp import GradScaler, autocast
import os
import datetime
from torch.utils.tensorboard import SummaryWriter

device = torch.device('cuda')

model_config = NanoGPTConfig()
train_config = TrainConfig()
print(train_config)
gpt = GPT(model_config).to(device)
print(f'Number of parameters:{count_parameters(gpt)}')

tokenizer = LlamaTokenizer.from_pretrained(train_config.tokenize_path)
tokenizer.add_special_tokens({'pad_token': '[PAD]'})

optimizer = torch.optim.AdamW(params=gpt.parameters(), lr=train_config.max_lr, betas=(0.9, 0.95))

loss_fn = CrossEntropyLoss()

scaler = GradScaler()

data_list = load_json_as_list(train_config.dataset_path)

dataset = CustomDataset(data_list)

data_loader = DataLoader(dataset, batch_size=train_config.batch_size, shuffle=True)

scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=len(data_loader), eta_min=train_config.min_lr)

print(f'total step:{len(data_loader)}')

start_step = 0
if train_config.load_model_path:
    checkpoint = torch.load(train_config.load_model_path)
    gpt.load_state_dict(checkpoint['model_state_dict'])
    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
    scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
    start_step = checkpoint['step']
    print(f"Resuming training from step {start_step}")

log_path = os.path.join(train_config.log_folder, train_config.name)
writer = SummaryWriter(log_path)

losses = []
for batch_idx, samples in enumerate(data_loader):
    if batch_idx < start_step:
        continue
    tokenized_sample = tokenizer.batch_encode_plus(samples,
                                                   max_length=model_config.block_size,
                                                   padding=True,
                                                   truncation=True,
                                                   return_tensors="pt")
    input_ids = tokenized_sample['input_ids']

    labels = input_ids.clone()
    labels[:, :-1] = labels[:, 1:]
    labels[:, -1] = tokenizer.pad_token_id

    input_ids, labels = input_ids.to(device), labels.to(device)

    with autocast():
        logits = gpt(input_ids)
        loss = loss_fn(logits.view(-1, logits.size(-1)), labels.view(-1))
        # loss = loss / train_config.accumulation_steps  # 将损失除以累积步数

    losses.append(loss.item())

    scaler.scale(loss).backward()

    # torch.nn.utils.clip_grad_norm_(gpt.parameters(), max_norm=train_config.clip_value)

    if (batch_idx + 1) % train_config.accumulation_steps == 0 or (batch_idx + 1) == len(data_loader):
        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad()

    scheduler.step()

    if batch_idx % train_config.log_freq == 0 or (batch_idx + 1) == len(data_loader):
        avg_loss = sum(losses) / len(losses)
        # avg_loss *= train_config.accumulation_steps
        writer.add_scalar('Training Loss', avg_loss, batch_idx)
        print(f'Step:{batch_idx} Loss:{avg_loss}')
        losses = []
        current_lr = scheduler.get_last_lr()[0]
        writer.add_scalar('Learning Rate', current_lr, batch_idx)

    if batch_idx % train_config.save_fred == 0 or (batch_idx + 1) == len(data_loader):
        if not os.path.exists(train_config.save_path):
            os.makedirs(train_config.save_path)
        file_name = train_config.name + '-ckpt-' + str(batch_idx) + '.pth'
        torch.save({
            'model_state_dict': gpt.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'scheduler_state_dict': scheduler.state_dict(),
            'step': batch_idx
        },
            os.path.join(train_config.save_path, file_name))
        print(f'{datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")} model saved')

writer.close()
这段代码有什么可以改进的地方？