从零实现 nanoGPT（一）：让莎士比亚变成模型能读懂的数据

系列说明 ：这是「从零实现 nanoGPT」的第一篇。我们不急着搭建模型，而是先解决一个基础但关键的问题：如何把一段原始文本，转换成"根据前面的字符预测下一个字符"的训练样本？

本文使用 Python、NumPy 和 PyTorch，从原始莎士比亚文本出发，实现字符编码、训练集与验证集划分，以及一个可复用的 get_batch() 函数。

一、模型怎样从一段文字中学习"接下来的字符"？

先不看代码，想象我们有一段文本：

复制代码

hello

假设任务是：已知前面出现的字符，预测下一个字符。

那么，这段短短的文本就可以提供多个预测任务：

已知的字符	希望预测的下一个字符
`h`	`e`
`he`	`l`
`hel`	`l`
`hell`	`o`

这就是本文要构造的数据形式。

不过，模型不能直接拿 Python 字符串进行数值计算。我们需要先把字符映射为整数，然后把连续的整数序列切成输入 x 和目标 y。

整条数据处理路线可以概括为：

python 复制代码

原始文本 input.txt
        ↓
统计字符，建立字符 ↔ 整数映射
        ↓
把文本编码为整数序列
        ↓
按顺序划分训练集 / 验证集
        ↓
保存 train.bin、val.bin、meta.pkl
        ↓
随机截取连续片段，构造 x 和 y
        ↓
检查形状、类型和错位关系

为了把原理说清楚，这次采用 nanoGPT 的字符级莎士比亚数据集，而不是 GPT-2 的 BPE 分词版本。在字符级任务里，每个不同的字符就是一个 token，例如 h、e、空格和换行符。

二、第一步：把字符转换成整数

2.1 建立词表

仍然从 hello 入手：

scss 复制代码

text = "hello"

chars = sorted(set(text))
vocab_size = len(chars)

print(chars)
print(vocab_size)

运行结果：

css 复制代码

['e', 'h', 'l', 'o']
4

这里的三个操作非常简单：

set(text)：提取文本中不重复的字符；
sorted(...)：给字符一个固定的顺序，保证多次运行时映射一致；
len(chars)：得到词表大小 vocab_size。

为什么 hello 有 5 个字符，词表却只有 4 个？因为字符 l 出现了两次，但在词表中只需要占一个位置。

2.2 构造双向映射

接下来给每个字符分配一个整数 ID：

scss 复制代码

stoi = {ch: i for i, ch in enumerate(chars)}
itos = {i: ch for i, ch in enumerate(chars)}

print(stoi)
print(itos)

输出：

css 复制代码

{'e': 0, 'h': 1, 'l': 2, 'o': 3}
{0: 'e', 1: 'h', 2: 'l', 3: 'o'}

变量名来自英文：

stoi：string to integer，将字符映射为整数；
itos：integer to string，将整数映射回字符。

注意，ID 只是编号。o 的编号为 3、e 的编号为 0，并不表示它们之间存在大小关系或者语义距离。

2.3 实现 encode() 和 decode()

有了双向映射，我们就可以写出编码和解码函数：

python 复制代码

def encode(s):
    return [stoi[ch] for ch in s]


def decode(ids):
    return "".join(itos[i] for i in ids)


ids = encode("hello")
print(ids)
print(decode(ids))

输出：

csharp 复制代码

[1, 0, 2, 2, 3]
hello

三、第二步：处理真正的莎士比亚文本

3.1 准备项目和数据

我们先建立这样一个目录：

arduino 复制代码

my-nanogpt/
├── data/
│   └── shakespeare_char/
│       └── input.txt
├── prepare.py
└── batch.py

在项目根目录执行：

bash 复制代码

pip install numpy torch

mkdir -p data/shakespeare_char

curl -L \
  https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt \
  -o data/shakespeare_char/input.txt

这是 Karpathy 提供的 Tiny Shakespeare 文本。打开 input.txt，开头可以看到这样的内容：

sql 复制代码

First Citizen:
Before we proceed any further, hear me speak.

All:
Speak, speak.

与之前的 hello 不同，这里面包含大写字母、小写字母、标点、空格和换行符。它们都是需要编码的字符。

3.2 为什么要划分训练集和验证集？

我们希望把一部分数据用于学习，把另一部分保留下来，用于后续检查模型在未参与训练的文本上的表现。

本文采用一个容易理解的方案：前 90% 作为训练集，后 10% 作为验证集。

ini 复制代码

n = int(len(text) * 0.9)

train_text = text[:n]
val_text = text[n:]

这里有一个细节：我们不在字符级别随机打乱整篇文章。

因为文本是有顺序的：如果先把所有字符打乱，原本相邻的字符关系就被破坏了，也就无法构造正确的下一字符预测任务。

按顺序划分不代表每次都从头到尾读取训练集。后面 get_batch() 仍然会在各自的数据集内部，随机选择连续片段的起始位置。

3.3 为什么保存成 .bin 和 .pkl？

把文本编码为整数后，我们需要保存两类信息：

文件	保存内容	用途
`train.bin`	训练集的整数 ID	供后续随机取样
`val.bin`	验证集的整数 ID	供后续验证
`meta.pkl`	`stoi`、`itos`、`vocab_size`	保存编码规则

这里使用 np.uint16 保存字符 ID：

ini 复制代码

train_ids = np.array(encode(train_text), dtype=np.uint16)
val_ids = np.array(encode(val_text), dtype=np.uint16)

uint16 是 16 位无符号整数，占 2 字节，能表示 0～65535 的整数。莎士比亚字符级词表只有 65 个字符，因此完全够用。

随后使用 tofile() 写入二进制文件：

arduino 复制代码

train_ids.tofile(DATA_DIR / "train.bin")
val_ids.tofile(DATA_DIR / "val.bin")

.bin 在这里就是按顺序保存的一串原始整数，并不包含字符映射规则。因此我们还要使用 pickle 保存 meta.pkl，以后才能把整数重新解码为文字。

3.4 完整实现：prepare.py

把前面的步骤合并起来，prepare.py 的完整代码如下：

python 复制代码

from pathlib import Path
import pickle
import numpy as np

DATA_DIR = Path(__file__).resolve().parent / "data" / "shakespeare_char"


def main():
    # 1. 读取原始文本
    text = (DATA_DIR / "input.txt").read_text(encoding="utf-8")

    # 2. 建立字符词表
    chars = sorted(set(text))
    vocab_size = len(chars)
    stoi = {ch: i for i, ch in enumerate(chars)}
    itos = {i: ch for i, ch in enumerate(chars)}

    def encode(s):
        return [stoi[ch] for ch in s]

    def decode(ids):
        return "".join(itos[i] for i in ids)

    # 3. 编码可逆性测试
    assert decode(encode(text)) == text

    # 4. 顺序划分训练集和验证集
    n = int(len(text) * 0.9)
    train_text = text[:n]
    val_text = text[n:]

    # 5. 编码成紧凑的整数数组
    train_ids = np.array(encode(train_text), dtype=np.uint16)
    val_ids = np.array(encode(val_text), dtype=np.uint16)

    # 6. 保存二进制数据
    train_ids.tofile(DATA_DIR / "train.bin")
    val_ids.tofile(DATA_DIR / "val.bin")

    # 7. 保存字符映射
    meta = {
        "vocab_size": vocab_size,
        "stoi": stoi,
        "itos": itos,
    }
    with open(DATA_DIR / "meta.pkl", "wb") as f:
        pickle.dump(meta, f)

    print(f"原始文本长度: {len(text):,} characters")
    print(f"词表大小: {vocab_size}")
    print(f"训练集: {len(train_ids):,} tokens")
    print(f"验证集: {len(val_ids):,} tokens")
    print("数据准备完成！")


if __name__ == "__main__":
    main()

在项目根目录运行：

arduino 复制代码

python prepare.py

使用原始数据时，预期输出为：

makefile 复制代码

原始文本长度: 1,115,394 characters
词表大小: 65
训练集: 1,003,854 tokens
验证集: 111,540 tokens
数据准备完成！

此时目录中会多出三个文件：

python 复制代码

data/shakespeare_char/
├── input.txt
├── train.bin
├── val.bin
└── meta.pkl

到这里，我们已经解决了"如何把文本变成整数并保存"的问题，但还没有得到真正的训练样本。

接下来进入本文最重要的部分。

四、第三步：把整数序列变成输入 x 和目标 y

4.1 为什么要错开一个位置？

先回到我们的 hello：

复制代码

原始片段： h  e  l  l  o
输入 x：   h  e  l  l
目标 y：   e  l  l  o

注意：y 不是另一段无关的文本，而是相对于 x 向后错开一个字符。

如果将 hello 编码为：

csharp 复制代码

[1, 0, 2, 2, 3]

则有：

ini 复制代码

ids = encode("hello")

x = ids[:4]      # [1, 0, 2, 2]
y = ids[1:5]     # [0, 2, 2, 3]

我们可以把每个位置理解成一道预测题：

预测位置	可以看到的输入	对应的目标
0	`h`	`e`
1	`he`	`l`
2	`hel`	`l`
3	`hell`	`o`

这里有一个十分重要的结论：

长度为 4 的输入序列，不只提供一道预测题，而是可以在 4 个位置上分别提供预测目标。

不过，能否同时计算这些预测任务，还取决于后续模型是否正确限制了每个位置所能访问的信息。在实现因果注意力时，第 0 个位置只能看 h，第 1 个位置只能看 he，以此类推；如果让前面的预测位置偷看到未来的字符，预测任务就失去意义了。

我们在第一篇只负责构造正确的 x 和 y，不展开注意力的实现。

4.2 block_size 和 batch_size 到底是什么？

在 nanoGPT 中，我们会经常遇到两个参数：

ini 复制代码

block_size = 8
batch_size = 4

它们分别控制：

block_size：每一条输入序列包含多少个连续 token；
batch_size：一次取出多少条这样的序列。

例如，block_size = 8、batch_size = 4，意味着我们一次取出 4 条样本，每条样本包含 8 个输入 token，以及对应的 8 个目标 token。

因此张量形状应该是：

ini 复制代码

x.shape = (4, 8)
y.shape = (4, 8)

一般记作：

$X,Y ∈ ZB×T X,Y \in \mathbb{Z}^{B\times T}$ X,Y ∈ ZB×T

其中 B = batch_size，T = block_size。这是整数 ID 的数组，还不是词向量。

如果每条样本都从随机位置取出，那么一批样本可以写成：

ini 复制代码

X = [样本 1 的 8 个字符 ID]
    [样本 2 的 8 个字符 ID]
    [样本 3 的 8 个字符 ID]
    [样本 4 的 8 个字符 ID]

Y = [样本 1 的下一个字符 ID 序列]
    [样本 2 的下一个字符 ID 序列]
    [样本 3 的下一个字符 ID 序列]
    [样本 4 的下一个字符 ID 序列]

4.3 如何随机截取一批样本？

假设整份编码数据存放在 data 中。一条样本的构造逻辑非常简单：

ini 复制代码

i = 10  # 假设从位置 10 开始

x = data[i : i + block_size]
y = data[i + 1 : i + block_size + 1]

x 从第 i 个位置开始，而 y 从第 i+1 个位置开始。两者长度完全相同，只是相差一个位置。

真正取一批样本时，我们需要先随机产生 batch_size 个起点：

ini 复制代码

ix = torch.randint(
    0,
    len(data) - block_size,
    (batch_size,)
)

torch.randint 的上界不包含在取值范围内，所以最后一个可能的起点是 len(data) - block_size - 1。这样不仅 x 不会越界，y 所需的"下一个字符"也一定存在。

然后把各条序列堆叠起来：

less 复制代码

x = torch.stack([
    torch.from_numpy(data[i:i + block_size].astype(np.int64))
    for i in ix.tolist()
])

这里有两点值得解释：

torch.stack() 会在新的维度上堆叠多条一维序列，因此最终得到形状为 (batch_size, block_size) 的二维 Tensor。
保存文件时使用 np.uint16 是为了节省空间；取样时转成 np.int64，对应的 PyTorch 类型是 torch.long，适合作为后续模型的 token 索引。

五、完整实现：batch.py

现在把上面的逻辑封装成可复用的 get_batch()。

python 复制代码

from pathlib import Path
import pickle
import numpy as np
import torch

DATA_DIR = Path(__file__).resolve().parent / "data" / "shakespeare_char"

# 加载词表信息
with open(DATA_DIR / "meta.pkl", "rb") as f:
    meta = pickle.load(f)

stoi = meta["stoi"]
itos = meta["itos"]


def encode(s):
    return [stoi[ch] for ch in s]


def decode(ids):
    return "".join(itos[int(i)] for i in ids)


def get_batch(split, batch_size=4, block_size=8):
    # 1. 检查数据集名称
    if split not in ("train", "val"):
        raise ValueError("split 必须是 'train' 或 'val'")

    # 2. 以内存映射方式读取二进制整数
    data = np.memmap(
        DATA_DIR / f"{split}.bin",
        dtype=np.uint16,
        mode="r",
    )

    if batch_size <= 0 or block_size <= 0:
        raise ValueError("batch_size 和 block_size 必须是正整数")
    if len(data) <= block_size:
        raise ValueError("数据长度必须大于 block_size")

    # 3. 随机选择 batch_size 个起点
    ix = torch.randint(
        0,
        len(data) - block_size,
        (batch_size,),
    )

    # 4. 构造输入序列 x
    x = torch.stack([
        torch.from_numpy(
            data[i:i + block_size].astype(np.int64)
        )
        for i in ix.tolist()
    ])

    # 5. 构造错开一个位置的目标序列 y
    y = torch.stack([
        torch.from_numpy(
            data[i + 1:i + block_size + 1].astype(np.int64)
        )
        for i in ix.tolist()
    ])

    return x, y

get_batch() 的返回值是两个 Tensor：

ini 复制代码

x, y = get_batch("train", batch_size=4, block_size=8)

六、回顾：第一篇究竟完成了什么？

现在，我们已经从一份原始文本，得到了一套可以重复使用的数据模块：

能够提取字符词表并建立双向映射；
decode(encode(text)) == text；
训练集和验证集按顺序分别保存；
能够按 batch_size、block_size 随机获取连续序列；
x.shape == y.shape == (batch_size, block_size)；
x.dtype == y.dtype == torch.long；
torch.equal(x[:, 1:], y[:, :-1])。

最终目录是：

python 复制代码

my-nanogpt/
├── data/
│   └── shakespeare_char/
│       ├── input.txt
│       ├── train.bin
│       ├── val.bin
│       └── meta.pkl
├── prepare.py
└── batch.py

第一篇到此结束。我们还没有实现模型，也没有写训练循环，这是刻意的：先确保数据正确，再讨论怎样让模型从数据中学习。

下一篇，我们将直接复用本篇的 get_batch()，从一个简单的字符预测模型开始。

从零实现 nanoGPT（一）：让莎士比亚变成模型能读懂的数据