"""第 19 章：亲手计算缩放点积注意力与因果掩码；PyTorch CPU。"""

from math import sqrt
import torch


def scaled_attention(query, key, value, allowed=None):
    scores = query @ key.transpose(-2, -1) / sqrt(query.shape[-1])
    if allowed is not None:
        if not allowed.any(dim=-1).all():
            raise ValueError("每个查询至少需要一个允许访问的位置")
        scores = scores.masked_fill(~allowed, float("-inf"))
    weights = scores.softmax(dim=-1)
    return weights @ value, weights


def example():
    query = torch.tensor([[1, 0], [0, 1], [1, 1]], dtype=torch.float64)[None, None]
    key = query.clone()
    value = torch.tensor([[10, 0], [0, 10], [10, 10]], dtype=torch.float64)[None, None]
    allowed = torch.ones(3, 3, dtype=torch.bool).tril()
    return query, key, value, allowed


def rows(tensor):
    return [[round(float(x), 4) for x in row] for row in tensor[0, 0]]


def main():
    torch.set_num_threads(1)
    query, key, value, allowed = example()
    full, full_weights = scaled_attention(query, key, value)
    causal, causal_weights = scaled_attention(query, key, value, allowed)
    print(f"Q/K/V 形状：{tuple(query.shape)}")
    print(f"未加掩码的权重：{rows(full_weights)}")
    print(f"第一行加权结果：{rows(full)[0]}")
    print(f"因果权重：{rows(causal_weights)}")
    print(f"因果输出：{rows(causal)}")
    changed_key, changed_value = key.clone(), value.clone()
    changed_key[:, :, 2] += 100
    changed_value[:, :, 2] -= 100
    changed, _ = scaled_attention(query, changed_key, changed_value, allowed)
    print(f"改变第三个位置后，前两个输出不变："
          f"{torch.equal(causal[:, :, :2], changed[:, :, :2])}")
    # 乘零不是屏蔽：分数 0 经指数函数后仍然贡献 1。
    scores = query @ key.transpose(-2, -1) / sqrt(2)
    wrong = (scores*allowed).softmax(-1)
    print(f"错误地将未来分数乘零，第一行未来权重合计：{wrong[0, 0, 0, 1:].sum():.4f}")


if __name__ == "__main__":
    main()