"""本课独立实验；在仓库根目录执行 .venv/bin/python lessons/16/experiment.py。"""
import sys, json, math
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
import torch
from torch import nn
from torch.nn import functional as F
from poetry_gpt.common import ROOT, DEFAULT_DATA, read_json, read_jsonl
from poetry_gpt.model import ModelConfig, PoetryGPT, CausalAttention, Block
from poetry_gpt.data import Tokenizer, SPECIAL, clean_record, keywords_for
from poetry_gpt.labs import show
torch.set_num_threads(2)
torch.manual_seed(26)

model = PoetryGPT(ModelConfig(12, context=8, width=16, heads=2, layers=1, dropout=0))
opt = torch.optim.AdamW(model.parameters(), lr=0.01)
x = torch.tensor([[1, 2, 3, 4]])
y = torch.tensor([[2, 3, 4, 5]])
for i in range(31):
    opt.zero_grad()
    _, loss = model(x, y)
    loss.backward()
    opt.step()
    if i % 10 == 0:
        show(f'第{i}次，同一道练习题误差', round(loss.item(), 5))
print('这只是验证程序能学会一道题，不是新诗质量证明。')