"""本课独立实验；在仓库根目录执行 .venv/bin/python lessons/25/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)

if not (DEFAULT_DATA / 'train.jsonl').exists():
    raise ValueError('先运行 poet prepare。')
sets = {n: {r['group'] for r in read_jsonl(DEFAULT_DATA / f'{n}.jsonl')} for n in ['train', 'val', 'test']}
show('学习与验证同组交集', len(sets['train'] & sets['val']))
show('学习与测试同组交集', len(sets['train'] & sets['test']))
assert not (sets['train'] & sets['val'] or sets['train'] & sets['test'] or sets['val'] & sets['test'])