"""本课独立实验；在仓库根目录执行 .venv/bin/python lessons/12/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=8, heads=2, layers=1, dropout=0)).eval()
a = torch.tensor([[1, 2, 3]])
b = torch.tensor([[1, 2, 9]])
first = model(a)[0]
second = model(b)[0]
show('前两个位置最大变化', (first[:, :2] - second[:, :2]).abs().max().item())
attention = CausalAttention(model.config).eval()
_, weights = attention(torch.randn(1, 3, 8), inspect=True)
show('第一个头权重', weights[0, 0].tolist())