从零开始打造一个简单 Transformer 语言模型 [译]

# 如果需要额外安装工具，运行这行代码就行 #
!pip install torch numpy

import math
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from torch.utils.data import Dataset, DataLoader

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super(MultiHeadAttention, self).__init__()
        assert d_model % num_heads == 0, "d_model 必须能被 num_heads 整除"
        
        self.d_model = d_model  # 嵌入向量的维度，比如每个词用多长的数字表示
        self.num_heads = num_heads  # 有几个“头”来分担注意力
        self.d_k = d_model // num_heads  # 每个头的维度
        
        # 为查询（Q）、键（K）、值（V）准备线性变换工具
        self.w_q = nn.Linear(d_model, d_model)
        self.w_k = nn.Linear(d_model, d_model)
        self.w_v = nn.Linear(d_model, d_model)
        self.w_o = nn.Linear(d_model, d_model)  # 最后的输出层

    def split_heads(self, x, batch_size):
        # 把输入从 (batch_size, seq_len, d_model) 变成 (batch_size, num_heads, seq_len, d_k)
        x = x.view(batch_size, -1, self.num_heads, self.d_k)
        return x.permute(0, 2, 1, 3)  # 调整顺序，方便计算

    def forward(self, q, k, v, mask=None):
        batch_size = q.size(0)
        
        # 把 Q、K、V 分成多个头
        q = self.split_heads(self.w_q(q), batch_size)  # 输出：(batch_size, num_heads, seq_len_q, d_k)
        k = self.split_heads(self.w_k(k), batch_size)
        v = self.split_heads(self.w_v(v), batch_size)
        
        # 计算注意力分数（“缩放点积注意力”）
        scores = torch.matmul(q, k.transpose(-2, -1))  # 形状：(batch_size, num_heads, seq_len_q, seq_len_k)
        scores = scores / math.sqrt(self.d_k)  # 缩放一下，防止数字太大

        # 如果有掩码，就用上（比如不想让模型偷看未来的词）
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)
        
        # 用 softmax 把分数变成概率
        attn_weights = F.softmax(scores, dim=-1)
        
        # 用概率“加权”值 V
        context = torch.matmul(attn_weights, v)  # 输出：(batch_size, num_heads, seq_len_q, d_k)

        # 把多个头的输出拼回去
        context = context.permute(0, 2, 1, 3).contiguous()
        context = context.view(batch_size, -1, self.d_model)
        
        # 最后通过一个线性层输出
        output = self.w_o(context)
        
        return output

class PositionwiseFeedForward(nn.Module):
    def __init__(self, d_model, d_ff):
        super(PositionwiseFeedForward, self).__init__()
        self.fc1 = nn.Linear(d_model, d_ff)  # 第一层
        self.fc2 = nn.Linear(d_ff, d_model)  # 第二层
        
    def forward(self, x):
        return self.fc2(F.relu(self.fc1(x)))  # 先放大，再缩小，中间加个激活

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_seq_length=5000):
        super(PositionalEncoding, self).__init__()
        
        # 创建一个位置编码表
        pe = torch.zeros(max_seq_length, d_model)
        position = torch.arange(0, max_seq_length, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        
        pe[:, 0::2] = torch.sin(position * div_term)  # 偶数位置用正弦
        pe[:, 1::2] = torch.cos(position * div_term)  # 奇数位置用余弦
        
        # 注册为缓冲区（不会被训练更新）
        self.register_buffer('pe', pe.unsqueeze(0))
        
    def forward(self, x):
        # 把位置信息加到输入上
        return x + self.pe[:, :x.size(1)]

class EncoderLayer(nn.Module):
    def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
        super(EncoderLayer, self).__init__()
        
        self.self_attn = MultiHeadAttention(d_model, num_heads)  # 自注意力
        self.feed_forward = PositionwiseFeedForward(d_model, d_ff)  # 前馈网络
        
        self.norm1 = nn.LayerNorm(d_model)  # 归一化
        self.norm2 = nn.LayerNorm(d_model)
        
        self.dropout1 = nn.Dropout(dropout)  # 随机丢弃，防止过拟合
        self.dropout2 = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        # 自注意力 + 残差连接 + 归一化
        attn_output = self.self_attn(x, x, x, mask)
        x = self.norm1(x + self.dropout1(attn_output))
        
        # 前馈网络 + 残差连接 + 归一化
        ff_output = self.feed_forward(x)
        x = self.norm2(x + self.dropout2(ff_output))
        
        return x

class DecoderLayer(nn.Module):
    def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
        super(DecoderLayer, self).__init__()
        
        self.self_attn = MultiHeadAttention(d_model, num_heads)  # 自注意力
        self.cross_attn = MultiHeadAttention(d_model, num_heads)  # 交叉注意力
        self.feed_forward = PositionwiseFeedForward(d_model, d_ff)  # 前馈网络
        
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.norm3 = nn.LayerNorm(d_model)
        
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)
        self.dropout3 = nn.Dropout(dropout)

    def forward(self, x, enc_output, src_mask=None, tgt_mask=None):
        # 自注意力
        self_attn_output = self.self_attn(x, x, x, tgt_mask)
        x = self.norm1(x + self.dropout1(self_attn_output))
        
        # 交叉注意力（看编码器输出）
        cross_attn_output = self.cross_attn(x, enc_output, enc_output, src_mask)
        x = self.norm2(x + self.dropout2(cross_attn_output))
        
        # 前馈网络
        ff_output = self.feed_forward(x)
        x = self.norm3(x + self.dropout3(ff_output))
        
        return x

class Encoder(nn.Module):
    def __init__(self, d_model, num_heads, d_ff, num_layers, dropout=0.1):
        super(Encoder, self).__init__()
        
        self.layers = nn.ModuleList([
            EncoderLayer(d_model, num_heads, d_ff, dropout)
            for _ in range(num_layers)
        ])
        
    def forward(self, x, mask=None):
        for layer in self.layers:
            x = layer(x, mask)
        return x

class Decoder(nn.Module):
    def __init__(self, d_model, num_heads, d_ff, num_layers, dropout=0.1):
        super(Decoder, self).__init__()
        
        self.layers = nn.ModuleList([
            DecoderLayer(d_model, num_heads, d_ff, dropout)
            for _ in range(num_layers)
        ])
        
    def forward(self, x, enc_output, src_mask=None, tgt_mask=None):
        for layer in self.layers:
            x = layer(x, enc_output, src_mask, tgt_mask)
        return x

class TransformerLM(nn.Module):
    def __init__(self, vocab_size, d_model, num_heads, d_ff, num_layers, dropout=0.1):
        super(TransformerLM, self).__init__()
        
        self.d_model = d_model
        
        # 词嵌入和位置编码
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.positional_encoding = PositionalEncoding(d_model)
        
        # 解码器堆栈
        self.decoder = Decoder(d_model, num_heads, d_ff, num_layers, dropout)
        
        # 输出层
        self.fc_out = nn.Linear(d_model, vocab_size)
        
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        # 嵌入 + 位置编码
        x = self.embedding(x) * math.sqrt(self.d_model)
        x = self.positional_encoding(x)
        x = self.dropout(x)
        
        # 通过解码器（仅用解码器架构）
        x = self.decoder(x, x, None, mask)
        
        # 输出词概率
        output = self.fc_out(x)
        
        return output

def create_causal_mask(seq_len):
    """创建因果掩码，防止模型看到未来的词"""
    mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1)
    return mask == 0  # True 的地方是可以关注的

class SimpleTokenizer:
    def __init__(self, texts=None):
        self.char_to_idx = {}
        self.idx_to_char = {}
        
        if texts:
            self.fit(texts)
    
    def fit(self, texts):
        # 收集所有独特字符
        unique_chars = set()
        for text in texts:
            unique_chars.update(text)
        
        # 创建字符到数字的映射
        self.char_to_idx = {char: idx for idx, char in enumerate(sorted(unique_chars))}
        self.idx_to_char = {idx: char for char, idx in self.char_to_idx.items()}
    
    def encode(self, text):
        return [self.char_to_idx[char] for char in text]  # 文本变数字
    
    def decode(self, indices):
        return ''.join([self.idx_to_char[idx] for idx in indices])  # 数字变文本
    
    @property
    def vocab_size(self):
        return len(self.char_to_idx)  # 词汇表大小

class TextDataset(Dataset):
    def __init__(self, texts, tokenizer, seq_length):
        self.tokenizer = tokenizer
        self.seq_length = seq_length  # 一次看多长的序列
        
        # 把所有文本变成数字并连起来
        self.data = []
        for text in texts:
            self.data.extend(tokenizer.encode(text))
        
        # 算能切出多少段
        self.num_sequences = max(0, len(self.data) - seq_length)

    def __len__(self):
        return self.num_sequences
    
    def __getitem__(self, idx):
        # 取一段输入和对应的目标
        input_seq = self.data[idx:idx + self.seq_length]
        target_seq = self.data[idx + 1:idx + self.seq_length + 1]
        
        return (
            torch.tensor(input_seq, dtype=torch.long),
            torch.tensor(target_seq, dtype=torch.long)
        )

def train_transformer_lm(model, dataloader, num_epochs, learning_rate, device):
    model.to(device)  # 搬到 GPU 或 CPU
    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)  # 优化器
    criterion = nn.CrossEntropyLoss()  # 损失函数
    
    for epoch in range(num_epochs):  # 训练几轮
        model.train()
        total_loss = 0

        for batch_idx, (inputs, targets) in enumerate(dataloader):
            inputs, targets = inputs.to(device), targets.to(device)
            
            # 创建因果掩码
            seq_len = inputs.size(1)
            causal_mask = create_causal_mask(seq_len).to(device)
            
            # 前向传播
            optimizer.zero_grad()
            outputs = model(inputs, causal_mask)

            # 计算损失
            loss = criterion(outputs.view(-1, outputs.size(-1)), targets.view(-1))
            
            # 反向传播，更新参数
            loss.backward()
            optimizer.step()
            
            total_loss += loss.item()

            if (batch_idx + 1) % 10 == 0:
                print(f'第 {epoch+1}/{num_epochs} 轮, 批次 {batch_idx+1}/{len(dataloader)}, '
                      f'损失: {total_loss / (batch_idx + 1):.4f}')
        
        print(f'第 {epoch+1}/{num_epochs} 轮, 平均损失: {total_loss / len(dataloader):.4f}')
    
    return model

def generate_text(model, tokenizer, start_text, max_length, temperature=1.0, device='cpu'):
    model.eval()  # 评估模式
    
    # 把起始文本变成数字
    input_seq = tokenizer.encode(start_text)
    input_tensor = torch.tensor([input_seq], dtype=torch.long).to(device)

    # 一个个生成新字符
    for _ in range(max_length):
        seq_len = input_tensor.size(1)
        causal_mask = create_causal_mask(seq_len).to(device)
        
        # 预测下一个
        with torch.no_grad():
            output = model(input_tensor, causal_mask)
        
        next_token_logits = output[0, -1, :] / temperature

        # 概率化并采样
        probabilities = F.softmax(next_token_logits, dim=-1)
        next_token = torch.multinomial(probabilities, 1).item()
        
        # 加到序列里
        input_tensor = torch.cat([
            input_tensor, 
            torch.tensor([[next_token]], dtype=torch.long).to(device)
        ], dim=1)

    # 解码成文本
    generated_tokens = input_tensor[0].tolist()
    generated_text = tokenizer.decode(generated_tokens)
    
    return generated_text

def main():
    # 一些示例文本
    texts = [
        "Hello, how are you doing today?",
        "Transformers are powerful neural network architectures.",
        "Language models can generate coherent text.",
        "PyTorch is a popular deep learning framework."
    ]

    # 设置参数
    seq_length = 20  # 一次看 20 个字符
    batch_size = 4   # 每次训练 4 组
    d_model = 64     # 词向量长度
    num_heads = 4    # 4 个注意力头
    d_ff = 256       # 前馈网络维度
    num_layers = 2   # 2 层解码器
    dropout = 0.1    # 丢弃率
    num_epochs = 10  # 训练 10 轮
    learning_rate = 0.001  # 学习速度

    # 选设备
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"使用设备: {device}")
    
    # 初始化分词器
    tokenizer = SimpleTokenizer(texts)
    print(f"词汇表大小: {tokenizer.vocab_size}")

    # 准备数据
    dataset = TextDataset(texts, tokenizer, seq_length)
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
    
    # 创建模型
    model = TransformerLM(
        vocab_size=tokenizer.vocab_size,
        d_model=d_model,
        num_heads=num_heads,
        d_ff=d_ff,
        num_layers=num_layers,
        dropout=dropout
    )

    # 看看模型长啥样
    print(model)
    print(f"参数数量: {sum(p.numel() for p in model.parameters())}")
    
    # 开始训练
    model = train_transformer_lm(model, dataloader, num_epochs, learning_rate, device)
    
    # 生成文本试试
    start_text = "Hello"
    generated_text = generate_text(model, tokenizer, start_text, max_length=50, device=device)
    print(f"生成文本:\n{generated_text}")