8：Transformer

import torch
import torch.nn as nn
from torch.nn import functional as F
import matplotlib.pyplot as plt

## 超参设置（每个参数都有讲究！）
batch_size = 16     # 批处理大小
block_size = 32     # 上下文长度
max_iters = 5000    # 训练轮数
learning_rate = 1e-3 # 学习率
device = 'cuda' if torch.cuda.is_available() else 'cpu'
n_embd = 64         # embedding维度
n_head = 4          # 多头注意力头数
n_layer = 4         # transformer层数
dropout = 0.0       # dropout率
torch.manual_seed(42) # 随机种子，保证可复现

# 读取莎士比亚文本
with open('input.txt', 'r', encoding='utf-8') as f:
    text = f.read()

print(f"数据集总字符数: {len(text):,}")
print(f"前300个字符预览:\n{text[:300]}")

# 构建词汇表
chars = sorted(list(set(text)))
vocab_size = len(chars)
print(f"词汇表: {''.join(chars)}")
print(f"词汇表大小: {vocab_size}")

# 简单tokenizer实现
stoi = {ch: i for i, ch in enumerate(chars)}
itos = {i: ch for i, ch in enumerate(chars)}
encode = lambda s: [stoi[c] for c in s]
decode = lambda l: ''.join([itos[i] for i in l])

# 测试编解码
print(f"编码测试: {encode('Hello World!')}")
print(f"解码测试: {decode(encode('Hello World!'))}")

def get_batch(split):
    """
    生成训练批次数据

    你需要实现：
    1. 从训练/验证集中随机采样
    2. 生成输入序列x和目标序列y
    3. 返回正确维度的tensor
    """
    ###################
    # 你的代码在这里
    ###################
    pass

class Head(nn.Module):
    """单头自注意力"""

    def __init__(self, head_size):
        super().__init__()
        # QKV投影矩阵
        self.key = nn.Linear(n_embd, head_size, bias=False)
        self.query = nn.Linear(n_embd, head_size, bias=False)
        self.value = nn.Linear(n_embd, head_size, bias=False)

        # 下三角mask矩阵（防止看到未来）
        self.register_buffer('tril', torch.tril(torch.ones(block_size, block_size)))
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        B, T, C = x.shape
        ###################
        # 注意力计算逻辑
        # 1. 计算Q, K, V
        # 2. 计算注意力分数
        # 3. 应用mask
        # 4. 加权求和
        ###################
        pass

class MultiHeadAttention(nn.Module):
    """多头注意力：并行处理多个注意力头"""

    def __init__(self, num_heads, head_size):
        super().__init__()
        self.heads = nn.ModuleList([Head(head_size) for _ in range(num_heads)])
        self.proj = nn.Linear(n_embd, n_embd)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        ###################
        # 多头并行计算与合并
        ###################
        pass

class FeedForward(nn.Module):
    """前馈神经网络：简单而强大"""
    ###################
    # 实现两层线性变换+激活
    ###################
    pass

class Block(nn.Module):
    """Transformer Block：communication + computation"""
    ###################
    # 组装完整的transformer块
    # LayerNorm + MultiHeadAttention + LayerNorm + FeedForward
    ###################
    pass

class BigramLanguageModel(nn.Module):
    """完整的语言模型"""

    def __init__(self):
        super().__init__()
        # embedding表
        self.token_embedding_table = nn.Embedding(vocab_size, n_embd)
        self.position_embedding_table = nn.Embedding(block_size, n_embd)

        # 堆叠transformer块
        self.blocks = nn.Sequential(*[Block(n_embd, n_head=n_head) for _ in range(n_layer)])
        self.ln_f = nn.LayerNorm(n_embd)
        self.lm_head = nn.Linear(n_embd, vocab_size)

    def forward(self, idx, targets=None):
        ###################
        # 完整前向传播
        # 1. token + position embedding
        # 2. 通过transformer块
        # 3. 输出预测
        # 4. 计算损失（如果有target）
        ###################
        pass

    def generate(self, idx, max_new_tokens):
        """文本生成：让AI说话！"""
        for _ in range(max_new_tokens):
            # 裁剪到block_size
            idx_cond = idx[:, -block_size:]
            # 预测下一个token
            logits, _ = self(idx_cond)
            logits = logits[:, -1, :]  # 只要最后一个时间步
            probs = F.softmax(logits, dim=-1)
            # 采样下一个token
            idx_next = torch.multinomial(probs, num_samples=1)
            idx = torch.cat((idx, idx_next), dim=1)
        return idx

@torch.no_grad()
def estimate_loss():
    """评估训练和验证损失"""
    out = {}
    model.eval()
    ###################
    # 计算平均损失
    ###################
    model.train()
    return out

# 实例化模型
model = BigramLanguageModel().to(device)
print(f"模型参数量: {sum(p.numel() for p in model.parameters())/1e6:.2f}M")

# 训练循环
optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
train_losses, val_losses = [], []

for iter in range(max_iters):
    ###################
    # 训练循环实现
    # 1. 获取batch
    # 2. 前向传播
    # 3. 计算损失
    # 4. 反向传播
    # 5. 参数更新
    # 6. 定期评估
    ###################
    pass

# 绘制损失曲线
plt.figure(figsize=(10, 5))
plt.plot(train_losses, label='Train Loss', color='#FF6B6B')
plt.plot(val_losses, label='Validation Loss', color='#4ECDC4')
plt.xlabel('Iteration')
plt.ylabel('Loss')
plt.title('Training Progress: Loss Curves')
plt.legend()
plt.grid(True, alpha=0.3)
plt.savefig('loss_plot.png', dpi=300, bbox_inches='tight')

# 生成莎士比亚风格文本
print("AI莎士比亚开始创作...")
context = torch.zeros((1, 1), dtype=torch.long, device=device)
generated_text = decode(model.generate(context, max_new_tokens=2000)[0].tolist())
print(generated_text)