# 配套文章：02-transformer-computation.md
# 本轮新增示例未执行；环境和运行方式见上级 GUIDE.md。

import math
import torch
from torch import nn

# 固定初始化种子方便在同一环境观察；它不会让随机模型获得语言能力。
torch.manual_seed(7)

class TinyDecoder(nn.Module):
    def __init__(self, vocab=50, width=8, heads=2, max_length=16):
        super().__init__()
        # 每个头平均分配隐藏维度，因此 width 必须能被 heads 整除。
        if width % heads:
            raise ValueError("隐藏维度必须能被头数整除")
        self.heads, self.head_dim = heads, width // heads
        self.token = nn.Embedding(vocab, width)
        self.position = nn.Embedding(max_length, width)
        self.norm1 = nn.LayerNorm(width)
        # 一次投影生成 Q/K/V，再沿最后一维切开；三部分参数互不相同。
        self.qkv = nn.Linear(width, 3 * width, bias=False)
        self.out = nn.Linear(width, width, bias=False)
        self.norm2 = nn.LayerNorm(width)
        self.ffn = nn.Sequential(nn.Linear(width, 4 * width), nn.GELU(),
                                 nn.Linear(4 * width, width))
        self.final_norm = nn.LayerNorm(width)
        self.lm_head = nn.Linear(width, vocab, bias=False)

    def forward(self, ids):
        batch, length = ids.shape
        # 本例只有等长、无填充输入；因果 mask 之外没有 padding mask。
        positions = torch.arange(length, device=ids.device)
        x = self.token(ids) + self.position(positions)
        q, k, v = self.qkv(self.norm1(x)).chunk(3, dim=-1)

        # [B,N,d] → [B,N,H,D] → [B,H,N,D]，让各头独立做矩阵乘法。
        def split_heads(tensor):
            return tensor.reshape(batch, length, self.heads, self.head_dim).transpose(1, 2)
        q, k, v = map(split_heads, (q, k, v))
        scores = q @ k.transpose(-2, -1) / math.sqrt(self.head_dim)
        # 上三角代表未来位置；设为负无穷后，softmax 权重变成零。
        future = torch.ones(length, length, dtype=torch.bool, device=ids.device).triu(1)
        weights = scores.masked_fill(future, float("-inf")).softmax(dim=-1)
        mixed = weights @ v
        # 合并各头，投影回原隐藏维度，再接残差与逐位置 FFN。
        mixed = mixed.transpose(1, 2).contiguous().reshape(batch, length, -1)
        x = x + self.out(mixed)
        x = x + self.ffn(self.norm2(x))
        logits = self.lm_head(self.final_norm(x))
        return logits, weights

# 延续第一篇的教学 ID；这里不使用真实模型的 Tokenizer。
ids = torch.tensor([[3, 8, 12, 5, 17, 42, 23, 6]], dtype=torch.long)
model = TinyDecoder().eval()
with torch.no_grad():
    logits, attention = model(ids)
    # 仅最后一个已知位置用于本次的下一个 Token 预测。
    probabilities = logits[0, -1].softmax(dim=-1)
    values, indices = probabilities.topk(3)
print("各位置词表分数形状：", tuple(logits.shape))  # [1,8,50]
print("各头注意力形状：", tuple(attention.shape))   # [1,2,8,8]
print("首头第一行：", attention[0, 0, 0])         # 只能读第一个位置
print("随机模型候选 ID 与概率：", list(zip(indices.tolist(), values.tolist())))