第 31 章：从零训练一个小型中文 GPT

class CharTokenizer:    """最简中文字符级分词器。"""    def __init__(self, text):        chars = sorted(list(set(text)))        self.stoi = {ch: i for i, ch in enumerate(chars)}  # string → int        self.itos = {i: ch for i, ch in enumerate(chars)}  # int → string        self.vocab_size = len(chars)    def encode(self, s):        return [self.stoi[ch] for ch in s if ch in self.stoi]    def decode(self, ids):        return ''.join([self.itos.get(i, '?') for i in ids])

def load_data(filepath):    with open(filepath, 'r', encoding='utf-8') as f:        text = f.read()    # 可选：剔除过短的行、过多的空白    text = text.replace('\r\n', '\n').replace('\r', '\n')    return text

import torchimport numpy as npdef get_batch(data, block_size, batch_size):    """从数据中随机取 batch_size 个长度为 block_size 的序列。"""    ix = torch.randint(0, len(data) - block_size, (batch_size,))    x = torch.stack([data[i:i + block_size] for i in ix])    y = torch.stack([data[i + 1:i + block_size + 1] for i in ix])    return x, y

import torch.nn as nnimport torch.nn.functional as Fimport mathclass MiniGPT(nn.Module):    def __init__(self, vocab_size, d_model=256, n_heads=8,                 n_layers=6, block_size=128, dropout=0.1):        super().__init__()        self.token_emb = nn.Embedding(vocab_size, d_model)        self.pos_emb = nn.Embedding(block_size, d_model)        self.blocks = nn.ModuleList([            TransformerBlock(d_model, n_heads, dropout)            for _ in range(n_layers)        ])        self.ln_f = nn.LayerNorm(d_model)        self.lm_head = nn.Linear(d_model, vocab_size)        self.block_size = block_size        # 初始化权重        self.apply(self._init_weights)    def _init_weights(self, module):        if isinstance(module, nn.Linear):            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)            if module.bias is not None:                torch.nn.init.zeros_(module.bias)        elif isinstance(module, nn.Embedding):            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)    def forward(self, idx):        B, T = idx.shape        tok_emb = self.token_emb(idx)                        # (B, T, d_model)        pos = torch.arange(0, T, device=idx.device)        pos_emb = self.pos_emb(pos)                          # (T, d_model)        x = tok_emb + pos_emb                                # 第24章：位置编码        for block in self.blocks:            x = block(x)                                     # Self-Attn + FFN        x = self.ln_f(x)                                     # 最终 LayerNorm        logits = self.lm_head(x)                             # (B, T, vocab_size)        return logits    @torch.no_grad()    def generate(self, idx, max_new_tokens, temperature=1.0):        """自回归生成——每次预测一个字符，然后把它接在输入后面。"""        for _ in range(max_new_tokens):            idx_cond = idx[:, -self.block_size:]  # 裁剪到 block_size            logits = self(idx_cond)               # (B, T, vocab_size)            logits = logits[:, -1, :] / temperature  # 只取最后一个位置            probs = F.softmax(logits, dim=-1)            idx_next = torch.multinomial(probs, num_samples=1)            idx = torch.cat((idx, idx_next), dim=1)        return idxclass TransformerBlock(nn.Module):    """第 24 章的核心组件，完整实现。"""    def __init__(self, d_model, n_heads, dropout=0.1):        super().__init__()        self.ln1 = nn.LayerNorm(d_model)        self.attn = nn.MultiheadAttention(            d_model, n_heads, dropout=dropout, batch_first=True)        self.ln2 = nn.LayerNorm(d_model)        self.ffn = nn.Sequential(            nn.Linear(d_model, 4 * d_model),            nn.GELU(),            nn.Linear(4 * d_model, d_model),            nn.Dropout(dropout),        )    def forward(self, x):        # causal mask：每个位置只能看它自己及之前的词        causal_mask = torch.triu(            torch.ones(x.size(1), x.size(1), device=x.device),            diagonal=1        ).bool()        # Self-Attention + 残差 + LayerNorm        attn_out, _ = self.attn(            x, x, x, attn_mask=causal_mask, need_weights=False)        x = self.ln1(x + attn_out)        # FFN + 残差 + LayerNorm        x = self.ln2(x + self.ffn(x))        return x

def train(model, data, tokenizer, epochs=10, batch_size=32,          block_size=128, lr=3e-4, device='cpu'):    model = model.to(device)    optimizer = torch.optim.AdamW(model.parameters(), lr=lr)    steps_per_epoch = len(data) // (batch_size * block_size)    for epoch in range(epochs):        model.train()        total_loss = 0        for step in range(steps_per_epoch):            x, y = get_batch(data, block_size, batch_size)            x, y = x.to(device), y.to(device)            logits = model(x)                         # (B, T, vocab_size)            loss = F.cross_entropy(                logits.view(-1, logits.size(-1)),                y.view(-1)            )            optimizer.zero_grad()            loss.backward()            # 梯度裁剪——防止某一步的梯度爆炸            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)            optimizer.step()            total_loss += loss.item()        avg_loss = total_loss / steps_per_epoch        perplexity = math.exp(avg_loss)        print(f"epoch {epoch + 1:>2d}/{epochs}  "              f"loss = {avg_loss:.4f}  ppl = {perplexity:.1f}")        # 每个 epoch 后生成一段样本文本        if epoch % 2 == 0 or epoch == epochs - 1:            model.eval()            start = torch.randint(0, len(data) - 10, (1,))            context = data[start:start + 10].unsqueeze(0).to(device)            gen = model.generate(context, max_new_tokens=80,                                 temperature=0.8)            print(f"  [生成样本] {tokenizer.decode(gen[0].tolist())}")            print()

# ---- 主流程 ----if __name__ == "__main__":    # 1. 加载数据    text = load_data("corpus.txt")          # 替换为你的文本路径    print(f"语料长度: {len(text):,} 字符")    # 2. 构建分词器并编码    tokenizer = CharTokenizer(text)    data = torch.tensor(tokenizer.encode(text), dtype=torch.long)    print(f"词汇表大小: {tokenizer.vocab_size}")    # 3. 创建模型    model = MiniGPT(        vocab_size=tokenizer.vocab_size,        d_model=256,        n_heads=8,        n_layers=6,        block_size=128,    )    n_params = sum(p.numel() for p in model.parameters())    print(f"模型参数: {n_params / 1e6:.1f}M")    # 4. 训练    device = 'cuda' if torch.cuda.is_available() else 'cpu'    print(f"训练设备: {device}")    train(model, data, tokenizer, epochs=10, device=device)    # 5. 保存    torch.save(model.state_dict(), "minigpt.pt")    print("模型已保存到 minigpt.pt")