"""第 25 章：相同矩阵乘法预算下，比模型宽度、数据量与颜色一致性。"""

import argparse
import json
from math import cos, exp, pi
from pathlib import Path
from statistics import mean, stdev
import torch
from torch import nn
from compute_budget import BUDGET, matrix_flops, parameters
from tiny_transformer import TinyTransformer
from small_language_model import PAD, corpus, vocabulary, make_windows, evaluate, corrupt_colors

SEEDS = [11, 22, 33]


def train(texts, validation, vocab, width, seed):
    torch.manual_seed(seed)
    model = TinyTransformer(len(vocab), width=width).double()
    assert sum(p.numel() for p in model.parameters()) == parameters(width)
    optimizer = torch.optim.AdamW(model.parameters(), lr=0.003, weight_decay=0.01)
    generator = torch.Generator().manual_seed(seed+1000)
    training = make_windows(texts, vocab, 32)
    x, y, _ = training
    cost = matrix_flops(width)
    steps = BUDGET//cost
    checkpoints = {max(1, round(steps*fraction)) for fraction in [0, .25, .5, .75, 1]}
    history, seen = [], set()
    for step in range(1, steps+1):
        model.train()
        batch = torch.randint(len(x), (16,), generator=generator)
        seen.update(batch.tolist())
        progress = step/steps
        multiplier = progress/.1 if progress <= .1 else .1+.9*(1+cos(pi*(progress-.1)/.9))/2
        for group in optimizer.param_groups:
            group["lr"] = .003*multiplier
        optimizer.zero_grad(set_to_none=True)
        logits = model(x[batch])
        loss = nn.functional.cross_entropy(logits.flatten(0, 1), y[batch].flatten(), ignore_index=PAD)
        if not torch.isfinite(loss):
            raise RuntimeError("非有限训练损失")
        loss.backward()
        norm = nn.utils.clip_grad_norm_(model.parameters(), 1.0, error_if_nonfinite=True)
        optimizer.step()
        if step in checkpoints:
            val_loss, accuracy = evaluate(model, validation)
            train_loss, _ = evaluate(model, training)
            history.append(dict(step=step, compute=step*cost, train_loss=train_loss,
                                val_loss=val_loss, accuracy=accuracy, grad_norm=float(norm)))
    # 预先规定用预算结束时的参数；不在每个种子中再挑选最好的一步。
    val_loss, accuracy = evaluate(model, validation)
    return model, dict(seed=seed, val_loss=val_loss, accuracy=accuracy, history=history,
                       unique_seen=len(seen), steps=steps, compute=steps*cost)


def run():
    torch.set_num_threads(1)
    torch.use_deterministic_algorithms(True)
    train_texts, validation, test = corpus()
    vocab = vocabulary(train_texts)
    assert vocabulary(train_texts[:48]) == vocab
    val_data, test_data = [make_windows(texts, vocab, 32) for texts in [validation, test]]
    configurations = [("48 条一致", 32, train_texts[:48]),
                      ("48 条加宽", 64, train_texts[:48]),
                      ("192 条一致", 32, train_texts),
                      ("192 条错配", 32, corrupt_colors(train_texts)),
                      ("48 条错配", 32, corrupt_colors(train_texts[:48]))]
    results, models = [], []
    for name, width, texts in configurations:
        trials, candidates = [], []
        for seed in SEEDS:
            model, trial = train(texts, val_data, vocab, width, seed)
            candidates.append(model)
            trials.append(trial)
        results.append(dict(name=name, width=width, documents=len(texts), trials=trials,
                            mean_val=mean(t["val_loss"] for t in trials)))
        models.append(candidates)
    selected = min(range(len(results)), key=lambda index: results[index]["mean_val"])
    # 只有选定配置后才访问测试评分；不根据测试结果调整配置或挑种子。
    test_scores = [evaluate(model, test_data) for model in models[selected]]
    return dict(budget=BUDGET, seeds=SEEDS, results=results, selected=selected,
                test_scores=test_scores, vocab_size=len(vocab))


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--out", type=Path, help="可选：保存全部实验指标 JSON")
    args = parser.parse_args()
    report = run()
    print(f"预算上限={BUDGET:,} 矩阵 FLOPs；种子={SEEDS}；验证/测试各 32 条")
    for result in report["results"]:
        trials = result["trials"]
        losses, accuracies = [[trial[key] for trial in trials] for key in ["val_loss", "accuracy"]]
        print(f"{result['name']}：参数={parameters(result['width']):,}，更新={trials[0]['steps']}，"
              f"验证损失={mean(losses):.3f}±{stdev(losses):.3f}，颜色={mean(accuracies):.1%}")
        print("  各种子验证损失：" + ", ".join(f"{loss:.3f}" for loss in losses))
    base = report["results"][0]["trials"]
    wider = report["results"][1]["trials"]
    clean = report["results"][2]["trials"]
    noisy = report["results"][3]["trials"]
    small_noisy = report["results"][4]["trials"]
    for label, left, right in [("加宽−基线", wider, base), ("增量−基线", clean, base),
                               ("192 条：错配−一致", noisy, clean),
                               ("48 条：错配−一致", small_noisy, base)]:
        differences = [a["val_loss"]-b["val_loss"] for a, b in zip(left, right)]
        print(label + "：" + ", ".join(f"{value:+.3f}" for value in differences))
    selected = report["results"][report["selected"]]
    losses, accuracies = zip(*report["test_scores"])
    print(f"按平均验证损失选择：{selected['name']}")
    print(f"测试损失={mean(losses):.3f}±{stdev(losses):.3f}；"
          f"各种子困惑度=" + ", ".join(f"{exp(loss):.3f}" for loss in losses))
    print("测试颜色正确数：" + ", ".join(f"{round(value*32)}/32" for value in accuracies))
    if args.out:
        args.out.parent.mkdir(parents=True, exist_ok=True)
        args.out.write_text(json.dumps(report, ensure_ascii=False, indent=2))
        print(f"实验记录：{args.out}")


if __name__ == "__main__":
    main()

1. 单步成本不同，同时更新 400 次会给加宽模型约 3.65 倍的矩阵训练计算；那会回答另一个问题。
2. 两组模型尺寸、批次与更新次数相同，只改变可抽取文本池的大小；不同文本数和累计阅读量是两回事。
3. 会把种子也变成被调选的超参数。不是绝对不能这样做，但必须计入选择过程与成本，不能拿幸运运行代表配置的通常表现。
4. 损失平均了全部有效位置，也衡量真实词元的概率大小；颜色准确率只看特定位置的最高分词元是否正确。
5. $100/32=3.125$ 个百分点。
6. 两组都尚未表现出良好的颜色关系预测，当前指标差异可能没有显现出数据中该信息的价值。
7. 不能。固定划分上的训练随机性，不等于数据来源、句式和任务变化带来的不确定性。
8. 例如五组都使用两倍原预算，继续使用预算结束时参数、三个相同种子和平均验证损失选择规则。报告全部运行及实际使用量，包括没有改善、颜色仍然错误或出现训练异常的情况。